// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Owns a live run end to end: builds the out directory, launches
// `forge train`, streams stdout/stderr to run.log + the console buffer,
// polls log.csv into the MetricsStore, and drives the Run state machine.
// Single-writer: all Run mutations flow through here and persist via
// RunStore immediately. Posts a local notification on finish/fail.
import Foundation
import UserNotifications

@Observable
@MainActor
final class RunSupervisor {
    let store: RunStore
    let metrics = MetricsStore()

    private(set) var activeRunID: UUID?
    private(set) var consoleLines: [String] = []
    private var runner: ProcessRunner?
    private var csvPoller: Task<Void, Never>?

    init(store: RunStore) {
        self.store = store
    }

    enum SupervisorError: LocalizedError {
        case busy, invalidConfig([String]), forgeNotConfigured
        var errorDescription: String? {
            switch self {
            case .busy:
                return "Un entraînement est déjà en cours — Forge Studio les exécute séquentiellement."
            case .invalidConfig(let errs):
                return "Config invalide : " + errs.joined(separator: " · ")
            case .forgeNotConfigured:
                return "Binaire forge non configuré (Réglages)."
            }
        }
    }

    private func transition(_ run: inout Run, to next: RunState) {
        guard RunState.transitions[run.state]?.contains(next) == true else { return }
        run.state = next
        store.upsert(run)
    }

    /// Public entry: creates the run immediately. If a training is live the
    /// run stays `queued` (persisted) and launches automatically when the
    /// active one exits — trainings never overlap.
    func start(config: ForgeConfig, datasetPath: String, name: String,
               resumeFrom: String? = nil) async throws {
        let errors = config.validationErrors
        guard errors.isEmpty else { throw SupervisorError.invalidConfig(errors) }
        guard ForgeBinaryLocator.savedBinaryURL != nil else {
            throw SupervisorError.forgeNotConfigured
        }

        let stamp = ISO8601DateFormatter().string(from: .now)
            .replacingOccurrences(of: ":", with: "-")
        let outDir = store.workspaceURL
            .appendingPathComponent("runs/\(name)-\(stamp)")
        try FileManager.default.createDirectory(at: outDir,
                                                withIntermediateDirectories: true)
        let run = Run(name: name, createdAt: .now, config: config,
                      datasetPath: datasetPath, outDirectory: outDir.path,
                      resumeCheckpoint: resumeFrom)
        try config.exportJSON().write(to: URL(fileURLWithPath: run.configPath))
        store.upsert(run)
        if activeRunID == nil {
            try await launch(run)
        } // else: stays queued; launchNextQueued() picks it up on exit
    }

    private func launchNextQueued() {
        guard activeRunID == nil,
              let next = store.runs.filter({ $0.state == .queued })
                  .min(by: { $0.createdAt < $1.createdAt }) else { return }
        Task {
            do {
                try await launch(next)
            } catch {
                var r = next
                r.failureReason = error.localizedDescription
                transition(&r, to: .launching)
                transition(&r, to: .failed)
            }
        }
    }

    private func launch(_ queued: Run) async throws {
        guard activeRunID == nil else { throw SupervisorError.busy }
        guard let forge = ForgeBinaryLocator.savedBinaryURL else {
            throw SupervisorError.forgeNotConfigured
        }
        var run = queued

        var args = ["train", "--config", run.configPath, "--data", run.datasetPath,
                    "--out", run.outDirectory]
        if let resume = run.resumeCheckpoint { args += ["--resume", resume] }

        transition(&run, to: .launching)
        activeRunID = run.id
        consoleLines.removeAll()
        await metrics.reset()

        let processRunner = ProcessRunner()
        runner = processRunner
        let (lines, exit) = try await processRunner.launch(
            executable: forge, arguments: args,
            currentDirectory: forge.deletingLastPathComponent())
        run.pid = await processRunner.pid
        transition(&run, to: .running)

        // Console + run.log + stdout events.
        let logURL = URL(fileURLWithPath: run.runLogPath)
        FileManager.default.createFile(atPath: logURL.path, contents: nil)
        let logHandle = try? FileHandle(forWritingTo: logURL)
        let runID = run.id
        Task { [weak self] in
            for await (line, isErr) in lines {
                logHandle?.write(Data((line + "\n").utf8))
                await MainActor.run {
                    guard let self else { return }
                    self.consoleLines.append(isErr ? "⚠︎ " + line : line)
                    if self.consoleLines.count > 5000 {
                        self.consoleLines.removeFirst(1000)
                    }
                }
            }
            try? logHandle?.close()
        }

        // CSV poller: 1 Hz incremental tail, updates run summary fields.
        let csvURL = URL(fileURLWithPath: run.logCSVPath)
        csvPoller = Task { [weak self] in
            while !Task.isCancelled {
                guard let self else { return }
                await self.metrics.ingestCSV(at: csvURL)
                let snap = await self.metrics.snapshot(maxPoints: 4, smoothing: 0)
                await MainActor.run {
                    guard var r = self.store.runs.first(where: { $0.id == runID })
                    else { return }
                    if let last = snap.lastPoint {
                        r.lastStep = last.step
                        r.lastTrainLoss = last.trainLoss
                    }
                    if let best = snap.bestVal {
                        r.bestValLoss = best.loss
                        r.bestValStep = best.step
                    }
                    self.store.upsert(r)
                }
                try? await Task.sleep(for: .seconds(1))
            }
        }

        // Exit watcher.
        Task { [weak self] in
            let status = await exit.value
            await MainActor.run {
                guard let self,
                      var r = self.store.runs.first(where: { $0.id == runID })
                else { return }
                self.csvPoller?.cancel()
                r.pid = nil
                r.exitCode = status.code
                if status.code == 0 {
                    self.transition(&r, to: .finished)
                } else if r.state == .finishing || status.wasSignaled {
                    self.transition(&r, to: .stopped)
                } else {
                    r.failureReason = "forge s'est terminé avec le code \(status.code) — voir run.log"
                    self.transition(&r, to: .failed)
                }
                self.activeRunID = nil
                self.runner = nil
                Self.notify(run: r)
                self.launchNextQueued()
            }
        }
    }

    // Local notification with the outcome; silently skipped when not
    // running from a bundle (bare `swift run`) or when denied.
    static func requestNotificationAuth() {
        guard Bundle.main.bundleIdentifier != nil else { return }
        UNUserNotificationCenter.current()
            .requestAuthorization(options: [.alert, .sound]) { _, _ in }
    }

    private static func notify(run: Run) {
        guard Bundle.main.bundleIdentifier != nil else { return }
        let content = UNMutableNotificationContent()
        switch run.state {
        case .finished:
            content.title = "Run terminé — \(run.name)"
            var body = ""
            if let loss = run.lastTrainLoss {
                body += String(format: "loss finale %.4f", loss)
            }
            if let best = run.bestValLoss {
                body += String(format: " · best val %.4f", best)
            }
            content.body = body.isEmpty ? "Entraînement achevé." : body
        case .failed:
            content.title = "Run échoué — \(run.name)"
            content.body = run.failureReason ?? "Voir run.log"
        default:
            return
        }
        UNUserNotificationCenter.current().add(
            UNNotificationRequest(identifier: run.id.uuidString,
                                  content: content, trigger: nil))
    }

    /// SIGTERM. Forge has no signal handler: the process dies immediately and
    /// the last ckpt_latest.bin is the recovery point (UI says so).
    func stop() async {
        guard let id = activeRunID,
              var run = store.runs.first(where: { $0.id == id }) else { return }
        transition(&run, to: .finishing)
        try? await runner?.terminate()
    }
}