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

import SwiftUI

/// Convert section: quantize, fuse adapters, and import HF repos — every job
/// tracked with progress and a before/after size story.
struct ConvertView: View {
    @Bindable var model: AppModel
    @State private var controller = ConvertController()

    var body: some View {
        ScrollView {
            VStack(alignment: .leading, spacing: ZyquoTheme.spacing16) {
                QuantizeCard(model: model, controller: controller)
                FuseCard(model: model, controller: controller)
                ImportCard(controller: controller)

                if !controller.jobs.isEmpty {
                    Text("Jobs")
                        .font(ZyquoTheme.headlineFont)
                    ForEach(controller.jobs.reversed()) { job in
                        ConvertJobRow(job: job)
                    }
                }
            }
            .padding(ZyquoTheme.spacing20)
        }
        .onChange(of: controller.finishedCount) {
            Task { await model.refresh() }
        }
    }
}

// MARK: - Controller

struct ConvertJob: Identifiable {
    let id = UUID()
    var title: String
    var stage: String
    var fraction: Double?
    var finished = false
    var error: String?
}

@Observable
@MainActor
final class ConvertController {
    var jobs: [ConvertJob] = []
    var finishedCount = 0

    func track(title: String, events: AsyncStream<ConversionEvent>) {
        var job = ConvertJob(title: title, stage: "starting")
        jobs.append(job)
        let index = jobs.count - 1
        Task {
            for await event in events {
                switch event {
                case .stage(let stage, let fraction):
                    job.stage = stage
                    job.fraction = fraction
                case .finished(let output):
                    job.stage = "done → \(output.lastPathComponent)"
                    job.finished = true
                    finishedCount += 1
                case .failed(let message):
                    job.error = message
                    job.finished = true
                }
                jobs[index] = job
            }
        }
    }
}

private struct ConvertJobRow: View {
    let job: ConvertJob

    var body: some View {
        HStack(spacing: ZyquoTheme.spacing12) {
            if job.error != nil {
                Image(systemName: "xmark.circle.fill").foregroundStyle(ZyquoTheme.danger)
            } else if job.finished {
                Image(systemName: "checkmark.circle.fill").foregroundStyle(ZyquoTheme.success)
            } else {
                ProgressView().controlSize(.small)
            }
            VStack(alignment: .leading, spacing: ZyquoTheme.spacing2) {
                Text(job.title)
                    .font(ZyquoTheme.bodyFont.weight(.medium))
                Text(job.error ?? job.stage)
                    .font(ZyquoTheme.captionFont)
                    .foregroundStyle(job.error != nil ? ZyquoTheme.danger : ZyquoTheme.textSecondary)
                    .lineLimit(2)
            }
            Spacer()
            if let fraction = job.fraction, !job.finished {
                ProgressView(value: fraction).frame(width: 100)
            }
        }
        .padding(ZyquoTheme.spacing12)
        .zyquoCard()
    }
}

// MARK: - Quantize

private struct QuantizeCard: View {
    @Bindable var model: AppModel
    let controller: ConvertController

    @State private var modelID: String?
    @State private var config = QuantConfig()

    private var source: LocalModel? { model.models.first { $0.id == modelID } }
    private var candidates: [LocalModel] {
        model.models.filter { $0.type == .llm && $0.quantization == nil }
    }

    var body: some View {
        VStack(alignment: .leading, spacing: ZyquoTheme.spacing12) {
            Label("Quantize", systemImage: "square.resize.down")
                .font(ZyquoTheme.headlineFont)
            Text("Shrink an fp16/bf16 model with affine quantization — runs natively in Swift.")
                .font(ZyquoTheme.captionFont)
                .foregroundStyle(ZyquoTheme.textSecondary)

            HStack(spacing: ZyquoTheme.spacing12) {
                Picker("Model", selection: $modelID) {
                    Text(candidates.isEmpty ? "No unquantized models" : "Choose…").tag(String?.none)
                    ForEach(candidates) { Text($0.name).tag(String?.some($0.id)) }
                }
                .frame(maxWidth: 280)

                Picker("Bits", selection: $config.bits) {
                    ForEach(QuantConfig.affineBits, id: \.self) { Text("\($0)-bit") }
                }
                .frame(width: 110)

                Picker("Group", selection: $config.groupSize) {
                    ForEach(QuantConfig.affineGroupSizes, id: \.self) { Text("g\($0)") }
                }
                .frame(width: 90)

                if let source, let params = source.parameterCount {
                    Text("\(ByteCountFormatter.string(fromByteCount: source.weightsSize, countStyle: .file)) → ~\(ByteCountFormatter.string(fromByteCount: config.predictedWeightBytes(parameterCount: params), countStyle: .file))")
                        .font(ZyquoTheme.monoSmallFont)
                        .foregroundStyle(ZyquoTheme.chartThroughput)
                }

                Spacer()

                Button("Quantize") { run() }
                    .buttonStyle(.borderedProminent)
                    .tint(ZyquoTheme.accent)
                    .disabled(source == nil)
            }
        }
        .padding(ZyquoTheme.spacing16)
        .frame(maxWidth: .infinity, alignment: .leading)
        .zyquoCard()
    }

    private func run() {
        guard let source else { return }
        let output = "\(source.name)-\(config.bits)bit"
        Task {
            let events = await ConversionService.shared.quantize(
                model: source, config: config, outputName: output)
            controller.track(title: "Quantize \(source.name) → \(config.label)", events: events)
        }
    }
}

// MARK: - Fuse

private struct FuseCard: View {
    @Bindable var model: AppModel
    let controller: ConvertController

    @State private var runID: UUID?
    @State private var dequantize = true

    private var completedRuns: [TrainingRun] {
        model.runs.filter { $0.state == .completed || !$0.checkpoints.isEmpty }
    }
    private var selectedRun: TrainingRun? { model.runs.first { $0.id == runID } }
    private var baseModel: LocalModel? {
        selectedRun.flatMap { run in model.models.first { $0.id == run.baseModelID } }
    }

    var body: some View {
        VStack(alignment: .leading, spacing: ZyquoTheme.spacing12) {
            Label("Fuse Adapter", systemImage: "arrow.triangle.merge")
                .font(ZyquoTheme.headlineFont)
            Text("Merge a trained adapter into its base to make a standalone model. On quantized bases, de-quantizing preserves small adapters (re-quantization rounds them away).")
                .font(ZyquoTheme.captionFont)
                .foregroundStyle(ZyquoTheme.textSecondary)

            HStack(spacing: ZyquoTheme.spacing12) {
                Picker("Run", selection: $runID) {
                    Text(completedRuns.isEmpty ? "No completed runs" : "Choose…").tag(UUID?.none)
                    ForEach(completedRuns) { Text($0.name).tag(UUID?.some($0.id)) }
                }
                .frame(maxWidth: 340)

                Toggle("De-quantize", isOn: $dequantize)
                    .help("Recommended for quantized bases — keeps the adapter's effect intact.")

                Spacer()

                Button("Fuse") { run() }
                    .buttonStyle(.borderedProminent)
                    .tint(ZyquoTheme.accent)
                    .disabled(selectedRun == nil || baseModel == nil)
            }
        }
        .padding(ZyquoTheme.spacing16)
        .frame(maxWidth: .infinity, alignment: .leading)
        .zyquoCard()
    }

    private func run() {
        guard let selectedRun, let baseModel else { return }
        Task {
            do {
                let adapters = await RunStore.shared.adaptersDirectory(for: selectedRun)
                let output = "\(baseModel.name)-fused-\(selectedRun.id.uuidString.prefix(6))"
                let events = try await ConversionService.shared.fuse(
                    baseModel: baseModel, adapterDirectory: adapters,
                    outputName: output, dequantize: dequantize)
                controller.track(title: "Fuse \(selectedRun.name)", events: events)
            } catch {
                controller.jobs.append(ConvertJob(
                    title: "Fuse \(selectedRun.name)", stage: "", fraction: nil,
                    finished: true, error: error.localizedDescription))
            }
        }
    }
}

// MARK: - Import from Hugging Face

private struct ImportCard: View {
    let controller: ConvertController

    @State private var repoID = ""
    @State private var quantize = true
    @State private var config = QuantConfig()

    var body: some View {
        VStack(alignment: .leading, spacing: ZyquoTheme.spacing12) {
            Label("Convert from Hugging Face", systemImage: "arrow.down.doc")
                .font(ZyquoTheme.headlineFont)
            Text("Convert any Hugging Face checkpoint to MLX format (handles PyTorch .bin too), optionally quantizing on the way.")
                .font(ZyquoTheme.captionFont)
                .foregroundStyle(ZyquoTheme.textSecondary)

            HStack(spacing: ZyquoTheme.spacing12) {
                TextField("org/repo (e.g. Qwen/Qwen3-0.6B)", text: $repoID)
                    .textFieldStyle(.roundedBorder)
                    .frame(maxWidth: 340)

                Toggle("Quantize", isOn: $quantize)
                if quantize {
                    Picker("Bits", selection: $config.bits) {
                        ForEach(QuantConfig.affineBits, id: \.self) { Text("\($0)-bit") }
                    }
                    .frame(width: 110)
                }

                Spacer()

                Button("Convert") { run() }
                    .buttonStyle(.borderedProminent)
                    .tint(ZyquoTheme.accent)
                    .disabled(repoID.isEmpty)
            }
        }
        .padding(ZyquoTheme.spacing16)
        .frame(maxWidth: .infinity, alignment: .leading)
        .zyquoCard()
    }

    private func run() {
        let repo = repoID.trimmingCharacters(in: .whitespaces)
        let output = repo.replacingOccurrences(of: "/", with: "--")
            + (quantize ? "-\(config.bits)bit" : "-mlx")
        Task {
            do {
                let events = try await ConversionService.shared.convert(
                    hfPath: repo, outputName: output, quantize: quantize ? config : nil)
                controller.track(title: "Convert \(repo)", events: events)
            } catch {
                controller.jobs.append(ConvertJob(
                    title: "Convert \(repo)", stage: "", fraction: nil,
                    finished: true, error: error.localizedDescription))
            }
        }
    }
}