// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Per-run checkpoints with Generate (streamed token-by-token) and Eval
// (loss/ppl) panels — thin drivers over `forge generate` / `forge eval`.
import SwiftUI

struct CheckpointsPanel: View {
    let run: Run
    @Binding var selected: String?
    @Environment(AppModel.self) private var app
    @State private var checkpoints: [Checkpoint] = []
    @State private var prompt = "Once upon a time"
    @State private var temp = 0.8
    @State private var topK = 40
    @State private var maxTokens = 200
    @State private var output = ""
    @State private var evalResult: String?
    @State private var busy = false
    @State private var errorText: String?

    var body: some View {
        VStack(alignment: .leading, spacing: 12) {
            Text("Checkpoints").font(.title3.weight(.semibold))
            if checkpoints.isEmpty {
                Text("Aucun checkpoint pour l'instant (écrits tous les \(run.config.train.checkpointEvery) steps).")
                    .foregroundStyle(.secondary)
            } else {
                Picker("Checkpoint", selection: $selected) {
                    ForEach(checkpoints) { ck in
                        Text("step \(ck.step) · \(ByteCountFormatter.string(fromByteCount: ck.sizeBytes, countStyle: .file))")
                            .tag(Optional(ck.path))
                    }
                }
                .frame(maxWidth: 420)

                GroupBox("Génération") {
                    VStack(alignment: .leading, spacing: 8) {
                        TextField("Prompt", text: $prompt)
                        HStack {
                            Text("temp")
                            Slider(value: $temp, in: 0.1...1.5).frame(width: 100)
                            Text(String(format: "%.2f", temp)).monospacedDigit()
                            Stepper("top-k \(topK)", value: $topK, in: 1...200, step: 10)
                            Stepper("tokens \(maxTokens)", value: $maxTokens,
                                    in: 20...1000, step: 20)
                            Spacer()
                            Button(busy ? "…" : "Générer") { generate() }
                                .disabled(busy || selected == nil)
                        }
                        .font(.callout)
                        if !output.isEmpty {
                            ScrollView {
                                Text(output)
                                    .font(.system(.body, design: .serif))
                                    .textSelection(.enabled)
                                    .frame(maxWidth: .infinity, alignment: .leading)
                            }
                            .frame(height: 160)
                            .padding(8)
                            .background(.quaternary.opacity(0.4),
                                        in: RoundedRectangle(cornerRadius: 6))
                        }
                    }
                }

                GroupBox("Évaluation") {
                    HStack {
                        Button(busy ? "…" : "Évaluer (50 batches)") { evaluate() }
                            .disabled(busy || selected == nil)
                        if let r = evalResult {
                            Text(r).font(.callout.monospacedDigit())
                        }
                        Spacer()
                    }
                }
            }
            if let e = errorText {
                Label(e, systemImage: "exclamationmark.triangle").foregroundStyle(.red)
            }
        }
        .onAppear { scan() }
    }

    private func scan() {
        let dir = URL(fileURLWithPath: run.outDirectory)
        let files = (try? FileManager.default.contentsOfDirectory(
            at: dir, includingPropertiesForKeys: [.fileSizeKey, .contentModificationDateKey]))
            ?? []
        checkpoints = files.compactMap { url in
            let name = url.lastPathComponent
            guard name.hasPrefix("ckpt_"), name.hasSuffix(".bin"),
                  let step = Int(name.dropFirst(5).dropLast(4)) else { return nil }
            let vals = try? url.resourceValues(forKeys: [.fileSizeKey,
                                                         .contentModificationDateKey])
            return Checkpoint(path: url.path, step: step,
                              sizeBytes: Int64(vals?.fileSize ?? 0),
                              modifiedAt: vals?.contentModificationDate ?? .now,
                              valLossAtStep: nil)
        }
        .sorted { $0.step > $1.step }
        if selected == nil { selected = checkpoints.first?.path }
    }

    private var tokenizerPath: String? {
        let dir = URL(fileURLWithPath: run.datasetPath)
        return (try? FileManager.default.contentsOfDirectory(
            at: dir, includingPropertiesForKeys: nil))?
            .first { $0.pathExtension == "model" }?.path
    }

    private func generate() {
        guard let ckpt = selected, let forge = ForgeBinaryLocator.savedBinaryURL,
              let tok = tokenizerPath else {
            errorText = "tokenizer introuvable dans le dossier du dataset"
            return
        }
        busy = true
        output = ""
        errorText = nil
        Task {
            defer { busy = false }
            do {
                let runner = ProcessRunner()
                let (lines, exit) = try await runner.launch(
                    executable: forge,
                    arguments: ["generate", "--checkpoint", ckpt,
                                "--tokenizer", tok, "--prompt", prompt,
                                "--temp", String(temp), "--top-k", String(topK),
                                "--max-tokens", String(maxTokens)],
                    currentDirectory: forge.deletingLastPathComponent())
                for await (line, isErr) in lines where !isErr {
                    output += line + "\n"
                }
                let status = await exit.value
                if status.code != 0 {
                    errorText = "forge generate a échoué (code \(status.code))"
                }
            } catch {
                errorText = error.localizedDescription
            }
        }
    }

    private func evaluate() {
        guard let ckpt = selected, let forge = ForgeBinaryLocator.savedBinaryURL
        else { return }
        let valBin = run.datasetPath + "/val.bin"
        busy = true
        evalResult = nil
        errorText = nil
        Task {
            defer { busy = false }
            do {
                let runner = ProcessRunner()
                let (lines, exit) = try await runner.launch(
                    executable: forge,
                    arguments: ["eval", "--checkpoint", ckpt, "--data", valBin,
                                "--batches", "50"],
                    currentDirectory: forge.deletingLastPathComponent())
                for await (line, _) in lines where line.hasPrefix("val loss") {
                    evalResult = line
                }
                let status = await exit.value
                if status.code != 0 {
                    errorText = "forge eval a échoué (code \(status.code))"
                }
            } catch {
                errorText = error.localizedDescription
            }
        }
    }
}