// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// Run dashboard: status strip, the loss chart (raw + bias-corrected EMA +
// val points, log/linear Y, best-val marker), and the console. Live data
// arrives via the supervisor's MetricsStore at ~1 Hz.
import Charts
import SwiftUI

struct RunDetailView: View {
    let run: Run
    @Environment(AppModel.self) private var app
    @State private var snapshot: MetricsStore.Snapshot?
    @State private var smoothing = 0.6
    @State private var logScale = true
    @State private var showConsole = false
    @State private var historicStore: MetricsStore?
    @State private var xState = ChartXState()
    @State private var checkpointSteps: [Int] = []
    @State private var selectedCheckpoint: String?

    private var isLive: Bool { app.supervisor.activeRunID == run.id }

    private func scanCheckpoints() {
        let dir = URL(fileURLWithPath: run.outDirectory)
        checkpointSteps = ((try? FileManager.default.contentsOfDirectory(
            at: dir, includingPropertiesForKeys: nil)) ?? [])
            .compactMap { url -> Int? in
                let n = url.lastPathComponent
                guard n.hasPrefix("ckpt_"), n.hasSuffix(".bin"),
                      let s = Int(n.dropFirst(5).dropLast(4)) else { return nil }
                return s
            }
            .sorted()
    }

    var body: some View {
        VStack(spacing: 0) {
            StatusStrip(run: run, snapshot: snapshot)
            Divider()
            ScrollView {
                VStack(alignment: .leading, spacing: 16) {
                    TrainingChartView(snapshot: snapshot, isLive: isLive,
                                      smoothing: $smoothing, logScale: $logScale,
                                      xState: xState,
                                      checkpointSteps: checkpointSteps,
                                      onCheckpointTap: { step in
                                          selectedCheckpoint = run.outDirectory
                                              + String(format: "/ckpt_%06d.bin", step)
                                      })
                    HStack(spacing: 24) {
                        secondaryChart(title: "Learning rate",
                                       series: snapshot?.lr ?? [], format: "%.2e")
                        secondaryChart(title: "Tokens/s",
                                       series: snapshot?.tokensPerSec ?? [], format: "%.0f")
                        secondaryChart(title: "Grad norm",
                                       series: snapshot?.gradNorm ?? [], format: "%.2f",
                                       threshold: run.config.train.gradClip)
                    }
                    if !run.state.isActive {
                        CheckpointsPanel(run: run, selected: $selectedCheckpoint)
                    }
                    DisclosureGroup("Console (\(app.supervisor.consoleLines.count) lignes)",
                                    isExpanded: $showConsole) {
                        ConsoleView(lines: isLive ? app.supervisor.consoleLines
                                                  : ["(run terminé — voir run.log)"])
                    }
                }
                .padding()
            }
        }
        .toolbar {
            ToolbarItem {
                Menu {
                    Button("Chart en PNG (2×)…") { exportPNG() }
                    Button("Métriques en CSV…") { exportCSV() }
                } label: {
                    Label("Exporter", systemImage: "square.and.arrow.up")
                }
                .disabled(snapshot == nil)
            }
            if isLive {
                ToolbarItem {
                    Button(role: .destructive) {
                        Task { await app.supervisor.stop() }
                    } label: {
                        Label("Stop", systemImage: "stop.fill")
                    }
                    .help("SIGTERM — forge ne checkpointe pas à l'arrêt ; reprise depuis ckpt_latest.bin")
                }
            }
        }
        .task(id: run.id) {
            scanCheckpoints()
            await refreshLoop()
        }
        .navigationTitle(run.name)
    }

    // Secondary charts share the main chart's X window and crosshair
    // (linked axes): zooming/hovering up top moves everything.
    private func secondaryChart(title: String, series: [Downsampler.XY],
                                format: String, threshold: Double? = nil) -> some View {
        VStack(alignment: .leading, spacing: 4) {
            HStack {
                Text(title).font(.caption.weight(.semibold))
                Spacer()
                if let h = xState.hoverX,
                   let near = series.min(by: { abs($0.x - h) < abs($1.x - h) }) {
                    Text(String(format: format, near.y))
                        .font(.caption.monospacedDigit())
                        .foregroundStyle(.teal)
                } else if let last = series.last {
                    Text(String(format: format, last.y))
                        .font(.caption.monospacedDigit())
                        .foregroundStyle(.secondary)
                }
            }
            Chart {
                ForEach(series, id: \.x) { p in
                    LineMark(x: .value("step", p.x), y: .value("v", p.y))
                        .foregroundStyle(.teal)
                }
                if let t = threshold, t > 0 {
                    RuleMark(y: .value("clip", t))
                        .lineStyle(.init(lineWidth: 1, dash: [3, 3]))
                        .foregroundStyle(.red.opacity(0.5))
                }
                if let h = xState.hoverX {
                    RuleMark(x: .value("hover", h))
                        .foregroundStyle(.secondary.opacity(0.5))
                        .lineStyle(.init(lineWidth: 1))
                }
            }
            .chartXAxis(.hidden)
            .modifier(LinkedScrollModifier(xState: xState))
            .frame(height: 110)
        }
    }

    private func exportPNG() {
        let chart = TrainingChartView(snapshot: snapshot, isLive: false,
                                      smoothing: $smoothing, logScale: $logScale,
                                      xState: ChartXState(),
                                      checkpointSteps: checkpointSteps)
            .frame(width: 1200, height: 500)
            .padding()
            .background(Color(nsColor: .windowBackgroundColor))
        let renderer = ImageRenderer(content: chart)
        renderer.scale = 2.0
        guard let image = renderer.nsImage,
              let tiff = image.tiffRepresentation,
              let png = NSBitmapImageRep(data: tiff)?
                  .representation(using: .png, properties: [:]) else { return }
        let panel = NSSavePanel()
        panel.nameFieldStringValue = "\(run.name)-loss.png"
        if panel.runModal() == .OK, let url = panel.url {
            try? png.write(to: url)
        }
    }

    private func exportCSV() {
        let panel = NSSavePanel()
        panel.nameFieldStringValue = "\(run.name)-metrics.csv"
        if panel.runModal() == .OK, let url = panel.url {
            try? FileManager.default.removeItem(at: url)
            try? FileManager.default.copyItem(
                at: URL(fileURLWithPath: run.logCSVPath), to: url)
        }
    }

    private func refreshLoop() async {
        if isLive {
            while !Task.isCancelled && app.supervisor.activeRunID == run.id {
                snapshot = await app.supervisor.metrics
                    .snapshot(maxPoints: 1200, smoothing: smoothing)
                try? await Task.sleep(for: .milliseconds(500))
            }
        }
        // Historical (or freshly finished): load the whole CSV once.
        let store = MetricsStore()
        historicStore = store
        await store.ingestCSV(at: URL(fileURLWithPath: run.logCSVPath))
        snapshot = await store.snapshot(maxPoints: 1200, smoothing: smoothing)
    }
}

struct StatusStrip: View {
    let run: Run
    let snapshot: MetricsStore.Snapshot?

    var body: some View {
        HStack(spacing: 20) {
            StateBadge(state: run.state)
            if let step = run.lastStep {
                VStack(alignment: .leading, spacing: 1) {
                    Text("step \(step + 1) / \(run.config.train.maxSteps)")
                        .font(.callout.monospacedDigit())
                    ProgressView(value: run.progress).frame(width: 140)
                }
            }
            if let last = snapshot?.lastPoint {
                metric("loss", String(format: "%.4f", last.trainLoss))
                metric("ppl", String(format: "%.1f", exp(last.trainLoss)))
                metric("tok/s", String(format: "%.0f", last.tokensPerSec))
                if let e = last.elapsedS {
                    metric("écoulé", Duration.seconds(e)
                        .formatted(.time(pattern: .hourMinuteSecond)))
                }
            }
            if let best = run.bestValLoss {
                metric("best val", String(format: "%.4f", best))
            }
            if run.state == .running, let growth = snapshot?.lastGrowth,
               Date.now.timeIntervalSince(growth) > 30 {
                Label("possiblement bloqué", systemImage: "exclamationmark.triangle")
                    .font(.caption)
                    .foregroundStyle(.orange)
                    .help("Aucune nouvelle métrique depuis \(Int(Date.now.timeIntervalSince(growth))) s — le run n'a pas été interrompu.")
            }
            Spacer()
            Text("\(run.config.model.paramCount.formatted(.number.notation(.compactName))) params")
                .foregroundStyle(.secondary)
        }
        .padding(.horizontal)
        .padding(.vertical, 10)
    }

    private func metric(_ label: String, _ value: String) -> some View {
        VStack(alignment: .leading, spacing: 1) {
            Text(label).font(.caption2).foregroundStyle(.secondary)
            Text(value).font(.callout.monospacedDigit())
        }
    }
}

struct ConsoleView: View {
    let lines: [String]

    var body: some View {
        ScrollViewReader { proxy in
            ScrollView {
                LazyVStack(alignment: .leading, spacing: 0) {
                    ForEach(Array(lines.enumerated()), id: \.offset) { i, line in
                        Text(line)
                            .font(.system(.caption, design: .monospaced))
                            .textSelection(.enabled)
                            .id(i)
                    }
                }
                .frame(maxWidth: .infinity, alignment: .leading)
            }
            .frame(height: 200)
            .background(.black.opacity(0.85))
            .clipShape(RoundedRectangle(cornerRadius: 6))
            .onChange(of: lines.count) {
                proxy.scrollTo(lines.count - 1, anchor: .bottom)
            }
        }
    }
}

extension View {
    @ViewBuilder
    func `if`<T: View>(_ condition: Bool, transform: (Self) -> T) -> some View {
        if condition { transform(self) } else { self }
    }
}