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

import Foundation

enum TrainingServiceError: LocalizedError {
    case memoryGate(verdict: MemoryVerdict, suggestions: [String])
    case qloraNeedsQuantizedBase
    case alreadyRunning(UUID)
    case nothingToResume

    var errorDescription: String? {
        switch self {
        case .memoryGate(_, let suggestions):
            "This configuration won't fit in memory on this Mac. Try: "
                + suggestions.joined(separator: " · ")
        case .qloraNeedsQuantizedBase:
            "QLoRA requires a quantized base model. Pick a quantized base (e.g. a -4bit model) or quantize it first in Convert."
        case .alreadyRunning(let id):
            "Run \(id.uuidString.prefix(8)) is already active."
        case .nothingToResume:
            "No saved adapter weights found to resume from."
        }
    }
}

/// Orchestrates fine-tuning runs (charter 3.B): builds the mlx-lm YAML
/// config, drives `zyquo_train.py` through `PythonRunner`, persists state and
/// metrics via `RunStore`, and exposes a typed event stream. Runs are
/// cancellable; resume is an honest warm start (docs/TRAINING-RESEARCH.md §5.3).
actor TrainingService {

    static let shared = TrainingService()

    private var activeRunID: UUID?
    private var activeStreamTask: Task<Void, Never>?

    // MARK: - Configure

    /// Create (and persist) a new run after memory-gating it.
    func createRun(
        name: String,
        baseModel: LocalModel,
        dataset: Dataset,
        method: FineTuneMethod,
        hyperParams: HyperParams
    ) async throws -> TrainingRun {
        // QLoRA = LoRA on a quantized base; LoRA on a quantized base IS QLoRA —
        // keep the user's mental model consistent with reality.
        if method == .qlora && baseModel.quantization == nil {
            throw TrainingServiceError.qloraNeedsQuantizedBase
        }

        let verdict = MemoryAdvisor.trainingVerdict(
            for: baseModel, method: method, params: hyperParams)
        if verdict == .wontFit {
            throw TrainingServiceError.memoryGate(
                verdict: verdict,
                suggestions: MemoryAdvisor.suggestions(
                    for: baseModel, method: method, params: hyperParams))
        }

        let id = UUID()
        let run = TrainingRun(
            id: id,
            name: name,
            baseModelID: baseModel.id,
            datasetID: dataset.id,
            method: method,
            hyperParams: hyperParams,
            state: .configured,
            directory: await RunStore.shared.directory(for: id),
            createdAt: .now,
            startedAt: nil,
            finishedAt: nil,
            completedIterations: 0,
            checkpoints: [],
            lastError: nil
        )
        try await RunStore.shared.create(run)
        try writeConfig(run: run, baseModel: baseModel, dataset: dataset, resumeFrom: nil)
        return run
    }

    // MARK: - Start / resume

    /// Start (or warm-start-resume) a run, streaming typed events.
    func start(run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resume: Bool = false)
        async throws -> AsyncStream<TrainingEvent>
    {
        if let activeRunID { throw TrainingServiceError.alreadyRunning(activeRunID) }

        var run = run
        var resumeAdapter: URL?
        if resume {
            guard let adapter = await RunStore.shared.latestAdapter(for: run) else {
                throw TrainingServiceError.nothingToResume
            }
            resumeAdapter = adapter
        }
        try writeConfig(run: run, baseModel: baseModel, dataset: dataset, resumeFrom: resumeAdapter)

        run.state = .running
        run.startedAt = run.startedAt ?? .now
        run.lastError = nil
        try await RunStore.shared.save(run)
        activeRunID = run.id

        let configPath = run.directory.appendingPathComponent("config.yaml").path
        let pythonStream = try await PythonRunner.shared.stream(
            script: "zyquo_train", arguments: ["--config", configPath])

        let (stream, continuation) = AsyncStream.makeStream(of: TrainingEvent.self)
        let runSnapshot = run

        activeStreamTask = Task {
            var current = runSnapshot
            do {
                for try await pythonEvent in pythonStream {
                    guard let event = MetricsStream.decode(pythonEvent) else { continue }
                    switch event {
                    case .metric(let metric):
                        await RunStore.shared.append(metric: metric, to: current)
                        current.completedIterations = max(
                            current.completedIterations, metric.iteration)
                    case .failed(let message):
                        current.lastError = message
                    default:
                        break
                    }
                    continuation.yield(event)
                }
                // A cancelled Task ends the stream without throwing — check
                // cancellation before declaring completion.
                if Task.isCancelled {
                    current.state = .cancelled
                } else {
                    current.state = current.lastError == nil ? .completed : .failed
                }
            } catch is CancellationError {
                current.state = .cancelled
            } catch {
                current.state = .failed
                current.lastError = error.localizedDescription
                continuation.yield(.failed(message: error.localizedDescription))
            }
            current.finishedAt = .now
            current.checkpoints = await RunStore.shared.checkpoints(for: current)
            try? await RunStore.shared.save(current)
            self.clearActive()
            continuation.finish()
        }

        return stream
    }

    /// Cancel the active run (terminates the Python process via the stream's
    /// onTermination; state persists as cancelled with checkpoints intact for
    /// warm-start resume).
    func cancel() {
        activeStreamTask?.cancel()
    }


    private func clearActive() {
        activeRunID = nil
        activeStreamTask = nil
    }

    // MARK: - Config generation

    /// Serialize the run into mlx-lm's YAML schema
    /// (exact keys/defaults per docs/TRAINING-RESEARCH.md §1.2; the
    /// config-only keys `lora_parameters` make YAML mandatory).
    private func writeConfig(
        run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resumeFrom: URL?
    ) throws {
        let hp = run.hyperParams
        let adapterPath = run.directory.appendingPathComponent("adapters").path

        // .qlora maps to upstream "lora" — QLoRA is implied by the quantized base.
        let fineTuneType = run.method == .qlora ? "lora" : run.method.rawValue

        var lines: [String] = [
            "# Generated by Zyquo MLX — run \(run.id.uuidString)",
            "model: \(yamlQuote(baseModel.directory.path))",
            "train: true",
            "data: \(yamlQuote(dataset.directory.path))",
            "fine_tune_type: \(fineTuneType)",
            "optimizer: \(hp.optimizer)",
            "mask_prompt: \(hp.maskPrompt)",
            "num_layers: \(hp.numLayers)",
            "batch_size: \(hp.batchSize)",
            "iters: \(hp.iterations)",
            "val_batches: \(hp.valBatches)",
            "learning_rate: \(hp.learningRate)",
            "steps_per_report: \(hp.stepsPerReport)",
            "steps_per_eval: \(hp.stepsPerEval)",
            "grad_accumulation_steps: \(hp.gradAccumulationSteps)",
            "save_every: \(hp.saveEvery)",
            "adapter_path: \(yamlQuote(adapterPath))",
            "max_seq_length: \(hp.maxSeqLength)",
            "grad_checkpoint: \(hp.gradCheckpoint)",
            "seed: \(hp.seed)",
        ]

        if let resumeFrom {
            lines.append("resume_adapter_file: \(yamlQuote(resumeFrom.path))")
        }

        if run.method != .full {
            lines.append("lora_parameters:")
            lines.append("  rank: \(hp.rank)")
            lines.append("  scale: \(hp.scale)")
            lines.append("  dropout: \(hp.dropout)")
            if let keys = hp.keys, !keys.isEmpty {
                lines.append("  keys: [\(keys.map(yamlQuote).joined(separator: ", "))]")
            }
        }

        try (lines.joined(separator: "\n") + "\n")
            .write(
                to: run.directory.appendingPathComponent("config.yaml"),
                atomically: true, encoding: .utf8)
    }

    private func yamlQuote(_ value: String) -> String {
        "\"" + value.replacingOccurrences(of: "\"", with: "\\\"") + "\""
    }
}