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

import Foundation

/// Persists training runs, their metrics history, and checkpoint inventory
/// under `~/Library/Application Support/ZyquoMLX/Runs/<run-id>/`.
///
/// Layout:
///   run.json          — the `TrainingRun` (state, config, our iteration counter)
///   metrics.jsonl     — append-only `TrainingMetric` history (chart replay)
///   adapters/         — mlx-lm adapter dir (adapters.safetensors, NNNNNNN_… checkpoints)
///   config.yaml       — the exact mlx-lm config the run used
actor RunStore {

    static let shared = RunStore()

    private let root: URL

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

    func directory(for id: UUID) -> URL {
        root.appendingPathComponent(id.uuidString, isDirectory: true)
    }

    func adaptersDirectory(for run: TrainingRun) -> URL {
        run.directory.appendingPathComponent("adapters", isDirectory: true)
    }

    // MARK: - CRUD

    func create(_ run: TrainingRun) throws {
        try FileManager.default.createDirectory(
            at: run.directory, withIntermediateDirectories: true)
        try save(run)
    }

    func save(_ run: TrainingRun) throws {
        try PersistenceService.saveJSON(run, to: run.directory.appendingPathComponent("run.json"))
    }

    func load(id: UUID) throws -> TrainingRun {
        try PersistenceService.loadJSON(
            TrainingRun.self, from: directory(for: id).appendingPathComponent("run.json"))
    }

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

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

    // MARK: - Metrics history

    func append(metric: TrainingMetric, to run: TrainingRun) {
        guard let data = try? JSONEncoder().encode(metric),
            let line = String(data: data, encoding: .utf8)
        else { return }
        let url = run.directory.appendingPathComponent("metrics.jsonl")
        if let handle = FileHandle(forWritingAtPath: url.path) {
            handle.seekToEndOfFile()
            handle.write(Data((line + "\n").utf8))
            try? handle.close()
        } else {
            try? (line + "\n").write(to: url, atomically: true, encoding: .utf8)
        }
    }

    func metricsHistory(for run: TrainingRun) -> [TrainingMetric] {
        guard
            let content = try? String(
                contentsOf: run.directory.appendingPathComponent("metrics.jsonl"),
                encoding: .utf8)
        else { return [] }
        let decoder = JSONDecoder()
        return content.split(separator: "\n").compactMap {
            $0.data(using: .utf8).flatMap { try? decoder.decode(TrainingMetric.self, from: $0) }
        }
    }

    // MARK: - Checkpoints

    /// Scan the adapter directory for numbered checkpoints
    /// (`{iter:07d}_adapters.safetensors` — docs/TRAINING-RESEARCH.md §2.1).
    func checkpoints(for run: TrainingRun) -> [Checkpoint] {
        let dir = adaptersDirectory(for: run)
        guard
            let files = try? FileManager.default.contentsOfDirectory(
                at: dir, includingPropertiesForKeys: [.creationDateKey])
        else { return [] }
        return files.compactMap { url -> Checkpoint? in
            let name = url.lastPathComponent
            guard name.hasSuffix("_adapters.safetensors"),
                let iteration = Int(name.prefix(7))
            else { return nil }
            let savedAt =
                (try? url.resourceValues(forKeys: [.creationDateKey]))?.creationDate ?? .now
            return Checkpoint(fileName: name, iteration: iteration, fileURL: url, savedAt: savedAt)
        }
        .sorted { $0.iteration < $1.iteration }
    }

    /// Latest usable adapter weights for resume/fuse/inference.
    func latestAdapter(for run: TrainingRun) -> URL? {
        let main = adaptersDirectory(for: run).appendingPathComponent("adapters.safetensors")
        if FileManager.default.fileExists(atPath: main.path) { return main }
        return checkpoints(for: run).last?.fileURL
    }
}