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

import SwiftUI

/// Evaluate section: base-vs-tuned side-by-side qualitative comparison with
/// per-side stats (held-out perplexity via the Python bridge uses the same
/// engine as `--test`; qualitative compare is the flagship workflow).
struct EvaluateView: View {
    @Bindable var model: AppModel
    @State private var controller = EvaluateController()

    var body: some View {
        if model.models.count < 2 {
            EmptyStateView(
                icon: "checkmark.seal",
                title: "Prove Your Fine-Tune",
                message: "Once you have a base model and a fine-tuned (fused) variant, compare them side by side on the same prompt — with tokens/sec and TTFT for each.")
        } else {
            VStack(spacing: 0) {
                HStack(spacing: ZyquoTheme.spacing12) {
                    Picker("Base", selection: $controller.baseID) {
                        Text("Choose…").tag(String?.none)
                        ForEach(model.models.filter { $0.type == .llm }) {
                            Text($0.name).tag(String?.some($0.id))
                        }
                    }
                    Picker("Candidate", selection: $controller.candidateID) {
                        Text("Choose…").tag(String?.none)
                        ForEach(model.models.filter { $0.type == .llm }) {
                            Text($0.name).tag(String?.some($0.id))
                        }
                    }
                }
                .padding(ZyquoTheme.spacing12)

                HStack(spacing: ZyquoTheme.spacing8) {
                    TextField("Prompt to compare on…", text: $controller.prompt)
                        .textFieldStyle(.roundedBorder)
                        .onSubmit { controller.run(models: model.models) }
                    Button(controller.isRunning ? "Running…" : "Compare") {
                        controller.run(models: model.models)
                    }
                    .buttonStyle(.borderedProminent)
                    .tint(ZyquoTheme.accent)
                    .disabled(controller.isRunning || controller.baseID == nil
                        || controller.candidateID == nil || controller.prompt.isEmpty)
                }
                .padding(.horizontal, ZyquoTheme.spacing12)
                .padding(.bottom, ZyquoTheme.spacing12)

                Rectangle().fill(ZyquoTheme.border).frame(height: ZyquoTheme.hairline)

                HStack(alignment: .top, spacing: 0) {
                    ResultColumn(title: "Base", output: controller.baseOutput, stats: controller.baseStats)
                    Rectangle().fill(ZyquoTheme.border).frame(width: ZyquoTheme.hairline)
                    ResultColumn(title: "Candidate", output: controller.candidateOutput, stats: controller.candidateStats)
                }

                if let error = controller.errorMessage {
                    Text(error)
                        .font(ZyquoTheme.captionFont)
                        .foregroundStyle(ZyquoTheme.danger)
                        .padding(ZyquoTheme.spacing8)
                }
            }
        }
    }
}

private struct ResultColumn: View {
    let title: String
    let output: String
    let stats: InferenceStats?

    var body: some View {
        VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {
            HStack {
                Text(title)
                    .font(ZyquoTheme.headlineFont)
                Spacer()
                if let stats {
                    Text(String(format: "%.1f tok/s · TTFT %.2fs", stats.tokensPerSecond, stats.ttft))
                        .font(ZyquoTheme.monoSmallFont)
                        .foregroundStyle(ZyquoTheme.chartThroughput)
                }
            }
            ScrollView {
                Text(output.isEmpty ? "—" : output)
                    .font(ZyquoTheme.bodyFont)
                    .foregroundStyle(ZyquoTheme.textPrimary)
                    .textSelection(.enabled)
                    .frame(maxWidth: .infinity, alignment: .leading)
            }
        }
        .padding(ZyquoTheme.spacing16)
        .frame(maxWidth: .infinity, maxHeight: .infinity, alignment: .topLeading)
    }
}

@Observable
@MainActor
final class EvaluateController {
    var baseID: String?
    var candidateID: String?
    var prompt = "Explain what a LoRA adapter is in one sentence."
    var baseOutput = ""
    var candidateOutput = ""
    var baseStats: InferenceStats?
    var candidateStats: InferenceStats?
    var isRunning = false
    var errorMessage: String?

    /// Sequential run (one model in memory at a time — verifiable unload
    /// between the two sides).
    func run(models: [LocalModel]) {
        guard let base = models.first(where: { $0.id == baseID }),
            let candidate = models.first(where: { $0.id == candidateID })
        else { return }
        isRunning = true
        errorMessage = nil
        baseOutput = ""
        candidateOutput = ""

        Task {
            do {
                (baseOutput, baseStats) = try await generate(model: base)
                (candidateOutput, candidateStats) = try await generate(model: candidate)
            } catch {
                errorMessage = error.localizedDescription
            }
            _ = try? await InferenceEngine.shared.unload()
            isRunning = false
        }
    }

    private func generate(model: LocalModel) async throws -> (String, InferenceStats?) {
        try await InferenceEngine.shared.load(model: model)
        var output = ""
        var stats: InferenceStats?
        var params = GenerationParams()
        params.maxTokens = 512
        let stream = try await InferenceEngine.shared.generate(
            messages: [.user(prompt)], params: params)
        for try await event in stream {
            switch event {
            case .chunk(let piece): output += piece
            case .finished(let s): stats = s
            }
        }
        return (output, stats)
    }
}