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

import Charts
import SwiftUI

/// Train section: run configurator + run list + live run detail with loss
/// curves, throughput, log console, checkpoints, and cancel/resume.
struct TrainView: View {
    @Bindable var model: AppModel
    @State private var controller = TrainingController()
    @State private var isConfiguring = false
    @State private var selectedRunID: UUID?

    var body: some View {
        Group {
            if let runID = selectedRunID ?? controller.activeRun?.id,
                let run = model.runs.first(where: { $0.id == runID }) ?? controller.activeRun
            {
                RunDetailView(
                    run: run, model: model, controller: controller,
                    back: {
                        selectedRunID = nil
                        Task { await model.refresh() }
                    })
            } else if model.runs.isEmpty {
                EmptyStateView(
                    icon: "flame",
                    title: canConfigure ? "Forge Your First Fine-Tune" : "No Training Runs Yet",
                    message: canConfigure
                        ? "Pick a base model, a dataset, and a method — Zyquo gates configurations against this Mac's memory and streams live loss curves while it trains."
                        : "Fine-tune any local model on your own data with LoRA, QLoRA, or full fine-tuning. Add a model and a dataset first, then start your first run.",
                    actionLabel: canConfigure ? "New Run" : nil,
                    action: canConfigure ? { isConfiguring = true } : nil)
            } else {
                runList
            }
        }
        .sheet(isPresented: $isConfiguring) {
            RunConfiguratorView(model: model, controller: controller) { run in
                selectedRunID = run.id
                Task { await model.refresh() }
            }
        }
    }

    private var canConfigure: Bool {
        !model.models.isEmpty && !model.datasets.isEmpty
    }

    private var runList: some View {
        ScrollView {
            LazyVStack(spacing: ZyquoTheme.spacing8) {
                HStack {
                    Spacer()
                    Button {
                        isConfiguring = true
                    } label: {
                        Label("New Run", systemImage: "plus")
                    }
                    .disabled(!canConfigure)
                }
                ForEach(model.runs) { run in
                    Button {
                        selectedRunID = run.id
                    } label: {
                        RunRow(run: run)
                    }
                    .buttonStyle(.plain)
                }
            }
            .padding(ZyquoTheme.spacing20)
        }
    }
}

// MARK: - Live training controller

@Observable
@MainActor
final class TrainingController {
    var activeRun: TrainingRun?
    var metrics: [TrainingMetric] = []
    var logLines: [String] = []
    var isTraining = false
    var lastError: String?

    func start(run: TrainingRun, baseModel: LocalModel, dataset: Dataset, resume: Bool = false) {
        metrics = resume ? metrics : []
        logLines.append(resume ? "Resuming (warm start from latest adapter)…" : "Starting run…")
        activeRun = run
        isTraining = true
        lastError = nil

        Task {
            do {
                let events = try await TrainingService.shared.start(
                    run: run, baseModel: baseModel, dataset: dataset, resume: resume)
                for await event in events {
                    switch event {
                    case .started(let model, let iterations):
                        logLines.append("Training \(model) for \(iterations) iterations")
                    case .metric(let metric):
                        metrics.append(metric)
                        if let loss = metric.trainLoss {
                            logLines.append(String(
                                format: "iter %d  train %.3f  %.0f tok/s  peak %.1f GB",
                                metric.iteration, loss, metric.tokensPerSecond ?? 0,
                                metric.peakMemoryGB ?? 0))
                        }
                        if let loss = metric.valLoss {
                            logLines.append(String(format: "iter %d  VAL %.3f", metric.iteration, loss))
                        }
                    case .checkpointSaved(let name, _):
                        logLines.append("checkpoint \(name)")
                    case .finished:
                        logLines.append("Training finished ✅")
                    case .failed(let message):
                        logLines.append("FAILED: \(message)")
                        lastError = message
                    }
                }
            } catch {
                lastError = error.localizedDescription
                logLines.append("FAILED: \(error.localizedDescription)")
            }
            isTraining = false
        }
    }

    func cancel() {
        Task { await TrainingService.shared.cancel() }
        logLines.append("Cancelling…")
    }
}

// MARK: - Configurator

struct RunConfiguratorView: View {
    @Bindable var model: AppModel
    let controller: TrainingController
    let onStart: (TrainingRun) -> Void

    @Environment(\.dismiss) private var dismiss
    @State private var baseModelID: String?
    @State private var datasetID: UUID?
    @State private var method: FineTuneMethod = .qlora
    @State private var hp = HyperParams()
    @State private var runName = ""
    @State private var startError: String?

    private var baseModel: LocalModel? { model.models.first { $0.id == baseModelID } }
    private var dataset: Dataset? { model.datasets.first { $0.id == datasetID } }

    private var verdict: MemoryVerdict? {
        baseModel.map { MemoryAdvisor.trainingVerdict(for: $0, method: method, params: hp) }
    }

    var body: some View {
        VStack(spacing: 0) {
            Text("New Training Run")
                .font(ZyquoTheme.titleFont)
                .padding(.top, ZyquoTheme.spacing20)

            Form {
                Section("Base & data") {
                    Picker("Base model", selection: $baseModelID) {
                        Text("Choose…").tag(String?.none)
                        ForEach(model.models.filter { $0.type == .llm }) { item in
                            Text("\(item.name)\(item.quantization != nil ? "  ·  quantized" : "")")
                                .tag(String?.some(item.id))
                        }
                    }
                    Picker("Dataset", selection: $datasetID) {
                        Text("Choose…").tag(UUID?.none)
                        ForEach(model.datasets) { dataset in
                            Text("\(dataset.name)  (\(dataset.trainCount) samples)")
                                .tag(UUID?.some(dataset.id))
                        }
                    }
                    Picker("Method", selection: $method) {
                        ForEach(FineTuneMethod.allCases, id: \.self) {
                            Text($0.displayName)
                        }
                    }
                    .help("QLoRA trains adapters on a quantized base — the memory-efficient default. Full fine-tuning updates every weight and needs far more memory.")
                }

                Section("Hyperparameters") {
                    Stepper("Iterations: \(hp.iterations)", value: $hp.iterations, in: 50...20000, step: 50)
                    Stepper("Batch size: \(hp.batchSize)", value: $hp.batchSize, in: 1...16)
                        .help("Larger batches train smoother but use more memory.")
                    Stepper("LoRA rank: \(hp.rank)", value: $hp.rank, in: 2...64, step: 2)
                        .help("Higher rank = more adapter capacity, slightly more memory.")
                    Stepper("Layers to adapt: \(hp.numLayers)", value: $hp.numLayers, in: 4...64, step: 4)
                        .help("LoRA is applied to the last N transformer layers (-1 = all).")
                    TextField("Learning rate", value: $hp.learningRate, format: .number)
                    Stepper("Max sequence length: \(hp.maxSeqLength)", value: $hp.maxSeqLength, in: 256...8192, step: 256)
                    Toggle("Gradient checkpointing", isOn: $hp.gradCheckpoint)
                        .help("Trades ~30% speed for a large activation-memory saving.")
                    Toggle("Mask prompt (loss on completions only)", isOn: $hp.maskPrompt)
                }

                if let baseModel, let verdict {
                    Section("Memory") {
                        LabeledContent("Estimated need") {
                            Text(ByteCountFormatter.string(
                                fromByteCount: MemoryAdvisor.trainingBytes(
                                    for: baseModel, method: method, params: hp),
                                countStyle: .memory))
                        }
                        LabeledContent("Verdict") { VerdictBadge(verdict: verdict) }
                        if verdict != .comfortable {
                            ForEach(
                                MemoryAdvisor.suggestions(for: baseModel, method: method, params: hp),
                                id: \.self
                            ) { suggestion in
                                Label(suggestion, systemImage: "lightbulb")
                                    .font(ZyquoTheme.captionFont)
                                    .foregroundStyle(ZyquoTheme.textSecondary)
                            }
                        }
                    }
                }

                if let startError {
                    Text(startError)
                        .font(ZyquoTheme.captionFont)
                        .foregroundStyle(ZyquoTheme.danger)
                }
            }
            .formStyle(.grouped)

            HStack {
                Button("Cancel") { dismiss() }
                Spacer()
                Button("Start Training") { start() }
                    .buttonStyle(.borderedProminent)
                    .tint(ZyquoTheme.accent)
                    .disabled(baseModel == nil || dataset == nil || verdict == .wontFit)
            }
            .padding(ZyquoTheme.spacing20)
        }
        .frame(width: 560, height: 640)
    }

    private func start() {
        guard let baseModel, let dataset else { return }
        Task {
            do {
                let name = runName.isEmpty
                    ? "\(baseModel.name) · \(method.displayName) · \(dataset.name)"
                    : runName
                let run = try await TrainingService.shared.createRun(
                    name: name, baseModel: baseModel, dataset: dataset,
                    method: method, hyperParams: hp)
                controller.start(run: run, baseModel: baseModel, dataset: dataset)
                dismiss()
                onStart(run)
            } catch {
                startError = error.localizedDescription
            }
        }
    }
}

// MARK: - Run detail

struct RunDetailView: View {
    let run: TrainingRun
    @Bindable var model: AppModel
    @Bindable var controller: TrainingController
    let back: () -> Void

    @State private var historicMetrics: [TrainingMetric] = []

    private var isActive: Bool { controller.activeRun?.id == run.id && controller.isTraining }
    private var metrics: [TrainingMetric] {
        isActive || controller.activeRun?.id == run.id ? controller.metrics : historicMetrics
    }

    var body: some View {
        ScrollView {
            VStack(alignment: .leading, spacing: ZyquoTheme.spacing16) {
                HStack {
                    Button {
                        back()
                    } label: {
                        Label("Runs", systemImage: "chevron.left")
                    }
                    .buttonStyle(.plain)
                    .foregroundStyle(ZyquoTheme.accent)

                    Spacer()

                    if isActive {
                        Button("Cancel Run") { controller.cancel() }
                            .tint(ZyquoTheme.danger)
                    } else if run.state == .cancelled || run.state == .failed {
                        Button("Resume (warm start)") { resume() }
                            .buttonStyle(.borderedProminent)
                            .tint(ZyquoTheme.accent)
                    }
                }

                VStack(alignment: .leading, spacing: ZyquoTheme.spacing4) {
                    HStack {
                        Text(run.name)
                            .font(ZyquoTheme.titleFont)
                        StatusPill(
                            text: isActive ? "Running" : run.state.rawValue.capitalized,
                            color: isActive ? ZyquoTheme.accent : ZyquoTheme.textSecondary)
                    }
                    Text("\(run.method.displayName) · \(run.hyperParams.iterations) iterations · rank \(run.hyperParams.rank) · batch \(run.hyperParams.batchSize)")
                        .font(ZyquoTheme.captionFont)
                        .foregroundStyle(ZyquoTheme.textSecondary)
                }

                // Loss chart
                lossChart

                // Throughput + memory strip
                if let latest = metrics.last(where: { $0.tokensPerSecond != nil }) {
                    HStack(spacing: ZyquoTheme.spacing20) {
                        stat("Tokens/sec", String(format: "%.0f", latest.tokensPerSecond ?? 0))
                        stat("Iteration", "\(metrics.last?.iteration ?? 0)/\(run.hyperParams.iterations)")
                        stat("Peak memory", String(format: "%.1f GB", latest.peakMemoryGB ?? 0))
                        if let tokens = latest.trainedTokens {
                            stat("Trained tokens", tokens.formatted())
                        }
                    }
                }

                // Log console
                VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {
                    Text("Console")
                        .font(ZyquoTheme.headlineFont)
                    ScrollViewReader { proxy in
                        ScrollView {
                            VStack(alignment: .leading, spacing: 2) {
                                ForEach(Array(consoleLines.enumerated()), id: \.offset) { index, line in
                                    Text(line)
                                        .font(ZyquoTheme.monoSmallFont)
                                        .foregroundStyle(ZyquoTheme.textSecondary)
                                        .id(index)
                                }
                            }
                            .frame(maxWidth: .infinity, alignment: .leading)
                            .padding(ZyquoTheme.spacing12)
                        }
                        .frame(height: 160)
                        .background(ZyquoTheme.surfaceSecondary)
                        .clipShape(RoundedRectangle(cornerRadius: ZyquoTheme.radiusSmall))
                        .onChange(of: consoleLines.count) {
                            proxy.scrollTo(consoleLines.count - 1, anchor: .bottom)
                        }
                    }
                }

                // Checkpoints
                if !checkpoints.isEmpty {
                    VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {
                        Text("Checkpoints")
                            .font(ZyquoTheme.headlineFont)
                        ForEach(checkpoints) { checkpoint in
                            HStack {
                                Image(systemName: "externaldrive")
                                    .foregroundStyle(ZyquoTheme.slate)
                                Text(checkpoint.fileName)
                                    .font(ZyquoTheme.monoSmallFont)
                                Spacer()
                                Text("iter \(checkpoint.iteration)")
                                    .font(ZyquoTheme.captionFont)
                                    .foregroundStyle(ZyquoTheme.textSecondary)
                            }
                            .padding(ZyquoTheme.spacing8)
                            .zyquoCard()
                        }
                    }
                }
            }
            .padding(ZyquoTheme.spacing20)
        }
        .task(id: run.id) {
            if controller.activeRun?.id != run.id {
                historicMetrics = await RunStore.shared.metricsHistory(for: run)
            }
        }
    }

    private var consoleLines: [String] {
        controller.activeRun?.id == run.id
            ? controller.logLines
            : metrics.compactMap { metric in
                if let loss = metric.trainLoss {
                    return String(format: "iter %d  train %.3f", metric.iteration, loss)
                }
                if let loss = metric.valLoss {
                    return String(format: "iter %d  VAL %.3f", metric.iteration, loss)
                }
                return nil
            }
    }

    private var checkpoints: [Checkpoint] {
        run.checkpoints
    }

    private var lossChart: some View {
        Chart {
            ForEach(Array(metrics.enumerated()), id: \.offset) { _, metric in
                if let loss = metric.trainLoss {
                    LineMark(
                        x: .value("Iteration", metric.iteration),
                        y: .value("Train loss", loss),
                        series: .value("Series", "Train"))
                        .foregroundStyle(ZyquoTheme.chartTrain)
                        .interpolationMethod(.monotone)
                }
            }
            ForEach(Array(metrics.enumerated()), id: \.offset) { _, metric in
                if let loss = metric.valLoss {
                    LineMark(
                        x: .value("Iteration", metric.iteration),
                        y: .value("Val loss", loss),
                        series: .value("Series", "Val"))
                        .foregroundStyle(ZyquoTheme.chartVal)
                    PointMark(
                        x: .value("Iteration", metric.iteration),
                        y: .value("Val loss", loss))
                        .foregroundStyle(ZyquoTheme.chartVal)
                        .symbolSize(30)
                }
            }
        }
        .chartLegend(.visible)
        .chartForegroundStyleScale([
            "Train": ZyquoTheme.chartTrain, "Val": ZyquoTheme.chartVal,
        ])
        .frame(height: 220)
        .padding(ZyquoTheme.spacing12)
        .zyquoCard()
        .overlay {
            if metrics.isEmpty {
                Text(isActive ? "Waiting for the first report…" : "No metrics recorded")
                    .font(ZyquoTheme.captionFont)
                    .foregroundStyle(ZyquoTheme.textTertiary)
            }
        }
    }

    private func stat(_ label: String, _ value: String) -> some View {
        VStack(alignment: .leading, spacing: ZyquoTheme.spacing2) {
            Text(label)
                .font(ZyquoTheme.captionFont)
                .foregroundStyle(ZyquoTheme.textTertiary)
            Text(value)
                .font(ZyquoTheme.monoFont)
                .foregroundStyle(ZyquoTheme.textPrimary)
        }
    }

    private func resume() {
        guard
            let baseModel = model.models.first(where: { $0.id == run.baseModelID }),
            let dataset = model.datasets.first(where: { $0.id == run.datasetID })
        else { return }
        controller.start(run: run, baseModel: baseModel, dataset: dataset, resume: true)
    }
}

struct RunRow: View {
    let run: TrainingRun

    var body: some View {
        HStack(spacing: ZyquoTheme.spacing12) {
            Image(systemName: "flame")
                .font(.system(size: 18, weight: .light))
                .foregroundStyle(stateColor)
                .frame(width: 32, height: 32)
                .background(stateColor.opacity(0.1))
                .clipShape(RoundedRectangle(cornerRadius: ZyquoTheme.radiusSmall))

            VStack(alignment: .leading, spacing: ZyquoTheme.spacing2) {
                Text(run.name)
                    .font(ZyquoTheme.bodyFont.weight(.medium))
                    .foregroundStyle(ZyquoTheme.textPrimary)
                HStack(spacing: ZyquoTheme.spacing8) {
                    Text(run.method.displayName)
                    Text("\(run.completedIterations)/\(run.hyperParams.iterations) iters")
                    if !run.checkpoints.isEmpty {
                        Text("\(run.checkpoints.count) checkpoints")
                    }
                }
                .font(ZyquoTheme.captionFont)
                .foregroundStyle(ZyquoTheme.textSecondary)
            }

            Spacer()

            StatusPill(text: run.state.rawValue.capitalized, color: stateColor)
            Text(run.createdAt.formatted(date: .abbreviated, time: .shortened))
                .font(ZyquoTheme.captionFont)
                .foregroundStyle(ZyquoTheme.textTertiary)
        }
        .padding(ZyquoTheme.spacing12)
        .zyquoCard()
    }

    private var stateColor: Color {
        switch run.state {
        case .running: ZyquoTheme.accent
        case .completed: ZyquoTheme.success
        case .paused, .configured: ZyquoTheme.slate
        case .cancelled: ZyquoTheme.warning
        case .failed: ZyquoTheme.danger
        }
    }
}