//
//  DatasetFormats.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation

/// A problem found in one dataset row, with a concrete fix
/// (charter 3.A: report malformed rows with fixes).
struct DatasetRowIssue: Identifiable, Codable, Hashable, Sendable {
    var id: Int { line }
    /// 1-based line number in the source file.
    var line: Int
    var problem: String
    var fix: String
}

/// One parsed sample, normalized for preview.
struct DatasetSample: Sendable {
    var messages: [(role: String, content: String)]
    var estimatedTokens: Int
}

/// Format detection + row validation for the JSONL shapes mlx-lm accepts
/// (docs/TRAINING-RESEARCH.md §4.2 — detection order: completions, chat, text).
enum DatasetFormats {

    static func detect(firstRow row: [String: Any]) -> DatasetFormat? {
        if row["prompt"] != nil && row["completion"] != nil { return .completions }
        if row["messages"] != nil { return .chat }
        if row["text"] != nil { return .text }
        return nil
    }

    /// Validate one row against a format. Returns an issue or nil when valid.
    static func validate(row: [String: Any], format: DatasetFormat, line: Int) -> DatasetRowIssue? {
        switch format {
        case .chat:
            guard let messages = row["messages"] as? [[String: Any]] else {
                return DatasetRowIssue(
                    line: line,
                    problem: "Missing or non-array \"messages\" key.",
                    fix: "Use {\"messages\": [{\"role\": \"user\", \"content\": \"…\"}, …]}.")
            }
            if messages.isEmpty {
                return DatasetRowIssue(
                    line: line, problem: "\"messages\" is empty.",
                    fix: "Provide at least a user and an assistant message.")
            }
            let validRoles: Set<String> = ["system", "user", "assistant", "tool"]
            for (i, message) in messages.enumerated() {
                guard let role = message["role"] as? String, validRoles.contains(role) else {
                    return DatasetRowIssue(
                        line: line,
                        problem: "Message \(i + 1) has a missing/invalid \"role\".",
                        fix: "Use one of: system, user, assistant, tool.")
                }
                guard message["content"] is String else {
                    return DatasetRowIssue(
                        line: line,
                        problem: "Message \(i + 1) (\(role)) has no string \"content\".",
                        fix: "Every message needs a \"content\" string.")
                }
            }
            if (messages.last?["role"] as? String) != "assistant" {
                return DatasetRowIssue(
                    line: line,
                    problem: "Conversation does not end with an assistant message.",
                    fix: "The final message must be the assistant reply the model should learn.")
            }
            return nil

        case .completions:
            guard row["prompt"] is String else {
                return DatasetRowIssue(
                    line: line, problem: "Missing string \"prompt\".",
                    fix: "Use {\"prompt\": \"…\", \"completion\": \"…\"}.")
            }
            guard row["completion"] is String else {
                return DatasetRowIssue(
                    line: line, problem: "Missing string \"completion\".",
                    fix: "Use {\"prompt\": \"…\", \"completion\": \"…\"}.")
            }
            return nil

        case .text:
            guard let text = row["text"] as? String, !text.isEmpty else {
                return DatasetRowIssue(
                    line: line, problem: "Missing or empty string \"text\".",
                    fix: "Use {\"text\": \"…\"} with non-empty content.")
            }
            return nil
        }
    }

    /// Normalize a valid row into a preview sample.
    static func sample(from row: [String: Any], format: DatasetFormat) -> DatasetSample {
        var messages: [(String, String)] = []
        switch format {
        case .chat:
            for message in row["messages"] as? [[String: Any]] ?? [] {
                messages.append(
                    (message["role"] as? String ?? "?", message["content"] as? String ?? ""))
            }
        case .completions:
            messages = [
                ("user", row["prompt"] as? String ?? ""),
                ("assistant", row["completion"] as? String ?? ""),
            ]
        case .text:
            messages = [("text", row["text"] as? String ?? "")]
        }
        let characters = messages.reduce(0) { $0 + $1.1.count }
        // ~4 chars/token heuristic; precise counts come from the tokenizer at
        // training time (trainer truncates over max-seq-length with a warning).
        return DatasetSample(messages: messages, estimatedTokens: max(1, characters / 4))
    }
}