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

import Foundation

enum DatasetImportError: LocalizedError {
    case emptyFile
    case undetectableFormat
    case allRowsInvalid(issues: [DatasetRowIssue])
    case tooFewSamples(count: Int)

    var errorDescription: String? {
        switch self {
        case .emptyFile:
            "The file contains no data rows."
        case .undetectableFormat:
            "Could not detect the dataset format. Rows must be JSON objects with \"messages\", \"prompt\"+\"completion\", or \"text\" keys."
        case .allRowsInvalid(let issues):
            "No valid rows found. First problem (line \(issues.first?.line ?? 1)): \(issues.first?.problem ?? "")"
        case .tooFewSamples(let count):
            "Only \(count) valid samples — at least 8 are needed for a meaningful train/valid split."
        }
    }
}

/// Result of a dataset import: what was written plus the validation report.
struct DatasetImportReport: Sendable {
    var dataset: Dataset
    var issues: [DatasetRowIssue]
    var skippedRows: Int
}

/// Imports, validates, splits, and previews training datasets
/// (charter 3.A). mlx-lm does not auto-split local JSONL
/// (docs/TRAINING-RESEARCH.md §4.5) — the split is ours.
actor DatasetService {

    static let shared = DatasetService()

    private let root: URL

    init(root: URL = PersistenceService.datasetsDirectory) {
        self.root = root
    }

    // MARK: - Import

    /// Import a JSONL file: validate every row, split into
    /// train/valid(/test), and write the dataset directory.
    func importJSONL(
        from sourceURL: URL,
        name: String,
        validFraction: Double = 0.1,
        testFraction: Double = 0.0,
        seed: UInt64 = 0
    ) throws -> DatasetImportReport {
        let content = try String(contentsOf: sourceURL, encoding: .utf8)
        let lines = content.split(separator: "\n", omittingEmptySubsequences: true)
        guard !lines.isEmpty else { throw DatasetImportError.emptyFile }

        var format: DatasetFormat?
        var validRows: [String] = []
        var issues: [DatasetRowIssue] = []
        var totalTokens = 0
        var maxTokens = 0

        for (index, line) in lines.enumerated() {
            let lineNumber = index + 1
            guard
                let data = line.data(using: .utf8),
                let row = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
            else {
                issues.append(
                    DatasetRowIssue(
                        line: lineNumber, problem: "Not a valid JSON object.",
                        fix: "Each line must be one complete JSON object (no trailing commas, no multi-line objects)."))
                continue
            }

            if format == nil {
                format = DatasetFormats.detect(firstRow: row)
                guard format != nil else { throw DatasetImportError.undetectableFormat }
            }

            if let issue = DatasetFormats.validate(row: row, format: format!, line: lineNumber) {
                issues.append(issue)
                continue
            }

            let sample = DatasetFormats.sample(from: row, format: format!)
            totalTokens += sample.estimatedTokens
            maxTokens = max(maxTokens, sample.estimatedTokens)
            validRows.append(String(line))
        }

        guard let detectedFormat = format else { throw DatasetImportError.undetectableFormat }
        guard !validRows.isEmpty else { throw DatasetImportError.allRowsInvalid(issues: issues) }
        guard validRows.count >= 8 else { throw DatasetImportError.tooFewSamples(count: validRows.count) }

        // Deterministic shuffle then split.
        var generator = SeededGenerator(seed: seed)
        validRows.shuffle(using: &generator)

        let testCount = Int(Double(validRows.count) * testFraction)
        let validCount = max(1, Int(Double(validRows.count) * validFraction))
        let trainCount = validRows.count - validCount - testCount

        let directory = root.appendingPathComponent(sanitize(name), isDirectory: true)
        try FileManager.default.createDirectory(at: directory, withIntermediateDirectories: true)

        let train = validRows[0..<trainCount]
        let valid = validRows[trainCount..<(trainCount + validCount)]
        let test = validRows[(trainCount + validCount)...]

        try write(rows: Array(train), to: directory.appendingPathComponent("train.jsonl"))
        try write(rows: Array(valid), to: directory.appendingPathComponent("valid.jsonl"))
        if !test.isEmpty {
            try write(rows: Array(test), to: directory.appendingPathComponent("test.jsonl"))
        }

        let dataset = Dataset(
            id: UUID(),
            name: name,
            directory: directory,
            format: detectedFormat,
            trainCount: trainCount,
            validCount: validCount,
            testCount: test.count,
            totalTokens: totalTokens,
            maxSequenceTokens: maxTokens,
            importedAt: .now
        )
        try PersistenceService.saveJSON(dataset, to: directory.appendingPathComponent("dataset.json"))

        return DatasetImportReport(
            dataset: dataset, issues: issues, skippedRows: issues.count)
    }

    // MARK: - Library

    func scan() throws -> [Dataset] {
        let fm = FileManager.default
        guard fm.fileExists(atPath: root.path) else { return [] }
        let entries = try fm.contentsOfDirectory(
            at: root, includingPropertiesForKeys: nil, options: .skipsHiddenFiles)
        return entries.compactMap {
            try? PersistenceService.loadJSON(
                Dataset.self, from: $0.appendingPathComponent("dataset.json"))
        }
        .sorted { $0.importedAt > $1.importedAt }
    }

    /// First N samples of a split, parsed for preview.
    func preview(dataset: Dataset, split: String = "train", count: Int = 8) throws -> [DatasetSample] {
        let url = dataset.directory.appendingPathComponent("\(split).jsonl")
        let content = try String(contentsOf: url, encoding: .utf8)
        return content.split(separator: "\n").prefix(count).compactMap { line in
            guard
                let data = line.data(using: .utf8),
                let row = (try? JSONSerialization.jsonObject(with: data)) as? [String: Any]
            else { return nil }
            return DatasetFormats.sample(from: row, format: dataset.format)
        }
    }

    func delete(_ dataset: Dataset) throws {
        try FileManager.default.removeItem(at: dataset.directory)
    }

    // MARK: - Helpers

    private func write(rows: [String], to url: URL) throws {
        try (rows.joined(separator: "\n") + "\n").write(to: url, atomically: true, encoding: .utf8)
    }

    private func sanitize(_ name: String) -> String {
        name.replacingOccurrences(of: "[^A-Za-z0-9._-]+", with: "-", options: .regularExpression)
    }
}

/// Deterministic RNG for reproducible splits (SplitMix64).
struct SeededGenerator: RandomNumberGenerator {
    private var state: UInt64
    init(seed: UInt64) { state = seed &+ 0x9E37_79B9_7F4A_7C15 }
    mutating func next() -> UInt64 {
        state &+= 0x9E37_79B9_7F4A_7C15
        var z = state
        z = (z ^ (z >> 30)) &* 0xBF58_476D_1CE4_E5B9
        z = (z ^ (z >> 27)) &* 0x94D0_49BB_1331_11EB
        return z ^ (z >> 31)
    }
}