// Author: Simon-Pierre Boucher — contact@spboucher.ai
//
// The main loss chart: raw train (low opacity) + bias-corrected EMA + val
// points, log/linear Y, best-val rule, hover crosshair with a full callout
// (step, train, EMA, val, ppl, lr, tok/s), pinch-zoom + horizontal pan with
// a "follow live" pill when zoomed during a running train, double-click to
// reset. Data arrives pre-downsampled from MetricsStore.
import Charts
import SwiftUI

// Shared X-window + crosshair state: the main chart drives it, secondary
// charts observe it, so zooming/hovering is synchronized everywhere.
@Observable
final class ChartXState {
    var visibleLength: Double?  // nil = fit all
    var scrollX: Double = 0
    var hoverX: Double?
}

struct TrainingChartView: View {
    let snapshot: MetricsStore.Snapshot?
    let isLive: Bool
    @Binding var smoothing: Double
    @Binding var logScale: Bool
    var xState: ChartXState
    var checkpointSteps: [Int] = []
    var onCheckpointTap: ((Int) -> Void)?

    @State private var followLive = true
    @State private var magnifyBase: Double?

    private var visibleLength: Double? {
        get { xState.visibleLength }
        nonmutating set { xState.visibleLength = newValue }
    }
    private var scrollX: Double {
        get { xState.scrollX }
        nonmutating set { xState.scrollX = newValue }
    }
    private var hoverX: Double? {
        get { xState.hoverX }
        nonmutating set { xState.hoverX = newValue }
    }

    private var maxX: Double { snapshot?.train.last?.x ?? 1 }
    private var minX: Double { snapshot?.train.first?.x ?? 0 }

    var body: some View {
        VStack(alignment: .leading, spacing: 8) {
            controls
            chart
                .frame(minHeight: 320)
        }
    }

    private var controls: some View {
        HStack {
            Text("Loss").font(.title3.weight(.semibold))
            if let h = hoverInfo {
                calloutText(h)
            }
            Spacer()
            if isLive && visibleLength != nil && !followLive {
                Button {
                    followLive = true
                } label: {
                    Label("suivre ⏵", systemImage: "forward.fill")
                        .font(.caption)
                }
                .buttonStyle(.borderedProminent)
                .controlSize(.small)
            }
            Toggle("log Y", isOn: $logScale).toggleStyle(.checkbox)
            Button("fit") {
                visibleLength = nil
                followLive = true
            }
            .controlSize(.small)
            HStack(spacing: 4) {
                Text("EMA")
                Slider(value: $smoothing, in: 0...0.99).frame(width: 110)
                Text(String(format: "%.2f", smoothing)).monospacedDigit()
            }
            .font(.caption)
        }
    }

    private var chart: some View {
        Chart {
            ForEach(snapshot?.train ?? [], id: \.x) { p in
                LineMark(x: .value("step", p.x), y: .value("loss", p.y),
                         series: .value("s", "train"))
                    .foregroundStyle(.blue.opacity(0.22))
            }
            ForEach(snapshot?.trainEMA ?? [], id: \.x) { p in
                LineMark(x: .value("step", p.x), y: .value("ema", p.y),
                         series: .value("s", "train EMA"))
                    .foregroundStyle(.blue)
            }
            ForEach(snapshot?.val ?? [], id: \.x) { p in
                LineMark(x: .value("step", p.x), y: .value("val", p.y),
                         series: .value("s", "val"))
                    .foregroundStyle(.orange)
                PointMark(x: .value("step", p.x), y: .value("val", p.y))
                    .foregroundStyle(.orange)
                    .symbolSize(24)
            }
            if let best = snapshot?.bestVal {
                RuleMark(y: .value("best", best.loss))
                    .foregroundStyle(.orange.opacity(0.4))
                    .lineStyle(.init(lineWidth: 1, dash: [4, 4]))
                    .annotation(position: .topTrailing) {
                        Text(String(format: "best val %.3f @ %d", best.loss, best.step))
                            .font(.caption2)
                            .foregroundStyle(.orange)
                    }
            }
            ForEach(checkpointSteps, id: \.self) { step in
                RuleMark(x: .value("ckpt", Double(step)))
                    .foregroundStyle(.green.opacity(0.25))
                    .lineStyle(.init(lineWidth: 1, dash: [2, 4]))
                    .annotation(position: .top, alignment: .leading) {
                        Image(systemName: "externaldrive.fill")
                            .font(.system(size: 7))
                            .foregroundStyle(.green.opacity(0.6))
                    }
            }
            if let x = hoverX {
                RuleMark(x: .value("hover", x))
                    .foregroundStyle(.secondary.opacity(0.5))
                    .lineStyle(.init(lineWidth: 1))
            }
        }
        .modifier(LogScaleModifier(enabled: logScale))
        .modifier(ScrollModifier(
            visibleLength: visibleLength,
            scrollX: Binding(get: { xState.scrollX },
                             set: { xState.scrollX = $0 })))
        .chartLegend(.visible)
        .chartOverlay { proxy in
            GeometryReader { geo in
                Color.clear
                    .contentShape(Rectangle())
                    .onContinuousHover { phase in
                        switch phase {
                        case .active(let pt):
                            let plot = geo[proxy.plotFrame!]
                            hoverX = proxy.value(atX: pt.x - plot.origin.x, as: Double.self)
                        case .ended:
                            hoverX = nil
                        }
                    }
                    .gesture(magnify)
                    .onTapGesture(count: 2) {
                        visibleLength = nil
                        followLive = true
                    }
                    .onTapGesture(count: 1) {
                        // Click near a checkpoint rule selects it (M4:
                        // "clickable → jumps to that checkpoint").
                        guard let x = hoverX,
                              let nearest = checkpointSteps
                                  .min(by: { abs(Double($0) - x) < abs(Double($1) - x) }),
                              abs(Double(nearest) - x) < max(4, (maxX - minX) * 0.01)
                        else { return }
                        onCheckpointTap?(nearest)
                    }
            }
        }
        .onChange(of: xState.scrollX) { old, new in
            // A user-initiated pan while live breaks auto-follow.
            if isLive, followLive, let len = visibleLength,
               abs(new - (maxX - len)) > len * 0.05 {
                followLive = false
            }
        }
        .onChange(of: snapshot?.count ?? 0) {
            if isLive, followLive, let len = visibleLength {
                scrollX = max(minX, maxX - len)
            }
        }
    }

    private var magnify: some Gesture {
        MagnifyGesture()
            .onChanged { g in
                let base = magnifyBase ?? (visibleLength ?? (maxX - minX))
                magnifyBase = base
                let span = maxX - minX
                let newLen = min(max(base / g.magnification, span * 0.01), span)
                visibleLength = newLen < span * 0.999 ? newLen : nil
                if followLive, let len = visibleLength {
                    scrollX = max(minX, maxX - len)
                }
            }
            .onEnded { _ in magnifyBase = nil }
    }

    // Nearest-point lookup for the callout (train series is x-sorted).
    private struct HoverInfo {
        var step: Int
        var train: Double?
        var ema: Double?
        var val: Double?
    }

    private var hoverInfo: HoverInfo? {
        guard let x = hoverX, let snap = snapshot, !snap.train.isEmpty else { return nil }
        func nearest(_ s: [Downsampler.XY]) -> Downsampler.XY? {
            guard !s.isEmpty else { return nil }
            var lo = 0, hi = s.count - 1
            while lo < hi {
                let mid = (lo + hi) / 2
                if s[mid].x < x { lo = mid + 1 } else { hi = mid }
            }
            if lo > 0, abs(s[lo - 1].x - x) < abs(s[lo].x - x) { lo -= 1 }
            return s[lo]
        }
        let t = nearest(snap.train)
        return HoverInfo(step: Int(t?.x ?? x),
                         train: t?.y,
                         ema: nearest(snap.trainEMA)?.y,
                         val: nearest(snap.val).flatMap {
                             abs($0.x - x) < max(8, (maxX - minX) * 0.02) ? $0.y : nil
                         })
    }

    private func calloutText(_ h: HoverInfo) -> some View {
        HStack(spacing: 10) {
            Text("step \(h.step)").foregroundStyle(.secondary)
            if let t = h.train {
                Text(String(format: "train %.4f", t)).foregroundStyle(.blue)
                Text(String(format: "ppl %.1f", exp(t))).foregroundStyle(.secondary)
            }
            if let e = h.ema {
                Text(String(format: "EMA %.4f", e)).foregroundStyle(.blue)
            }
            if let v = h.val {
                Text(String(format: "val %.4f", v)).foregroundStyle(.orange)
            }
        }
        .font(.caption.monospacedDigit())
        .padding(.horizontal, 8)
        .padding(.vertical, 3)
        .background(.quaternary.opacity(0.5), in: Capsule())
    }
}

// Conditional modifiers kept out of the Chart body so the result builder
// stays type-checkable.
private struct LogScaleModifier: ViewModifier {
    let enabled: Bool
    func body(content: Content) -> some View {
        if enabled { content.chartYScale(type: .log) } else { content }
    }
}

// Read-only link to the main chart's X window (secondary charts).
struct LinkedScrollModifier: ViewModifier {
    var xState: ChartXState
    func body(content: Content) -> some View {
        if let len = xState.visibleLength {
            content
                .chartScrollableAxes(.horizontal)
                .chartXVisibleDomain(length: len)
                .chartScrollPosition(x: Binding(get: { xState.scrollX },
                                                set: { _ in })) // main chart drives
        } else {
            content
        }
    }
}

private struct ScrollModifier: ViewModifier {
    let visibleLength: Double?
    @Binding var scrollX: Double
    func body(content: Content) -> some View {
        if let len = visibleLength {
            content
                .chartScrollableAxes(.horizontal)
                .chartXVisibleDomain(length: len)
                .chartScrollPosition(x: $scrollX)
        } else {
            content
        }
    }
}