//
//  CompareView.swift
//  Zyquo Local
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import SwiftUI

/// Compare mode: two models side by side, same prompt broadcast, independent
/// streaming and stats. RAM-gated: both models must fit together.
struct CompareView: View {
    @Environment(AppModel.self) private var app
    @State private var controller = CompareController()
    @State private var prompt = ""
    @FocusState private var focused: Bool

    var body: some View {
        VStack(spacing: 0) {
            HStack(spacing: ZyquoTheme.Spacing.m) {
                columnHeader(side: 0)
                Rectangle().fill(ZyquoTheme.border).frame(width: ZyquoTheme.hairline)
                columnHeader(side: 1)
            }
            .frame(height: ZyquoTheme.chatHeaderHeight)
            .padding(.horizontal, ZyquoTheme.Spacing.m)

            if let warning = controller.ramWarning {
                Text(warning)
                    .font(ZyquoTheme.caption)
                    .foregroundStyle(ZyquoTheme.warning)
                    .padding(.bottom, ZyquoTheme.Spacing.xxs)
            }

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

            HStack(spacing: 0) {
                column(side: 0)
                Rectangle().fill(ZyquoTheme.border).frame(width: ZyquoTheme.hairline)
                column(side: 1)
            }

            inputRow
        }
        .background(ZyquoTheme.background)
        .frame(minWidth: 900, minHeight: 560)
        .onDisappear {
            Task { await controller.teardown() }
        }
    }

    private func columnHeader(side: Int) -> some View {
        HStack {
            Picker(
                "Model",
                selection: Binding(
                    get: { controller.selection[side] },
                    set: { newValue in
                        controller.selection[side] = newValue
                        Task { await controller.loadSelection(app: app) }
                    }
                )
            ) {
                Text("Choose…").tag(String?.none)
                ForEach(app.store.models) { model in
                    Text(model.name).tag(String?.some(model.repoID))
                }
            }
            .labelsHidden()
            .frame(maxWidth: 280)

            switch controller.states[side] {
            case .loading:
                ProgressView().controlSize(.mini)
            case .ready:
                Circle().fill(ZyquoTheme.success).frame(width: 7, height: 7)
            default:
                EmptyView()
            }
            Spacer()
            if let stats = controller.stats[side] {
                Text(String(format: "⚡ %.1f tok/s · %d tok · %.1fs TTFT",
                            stats.tokensPerSecond, stats.generationTokenCount, stats.timeToFirstToken))
                    .font(ZyquoTheme.caption.monospacedDigit())
                    .foregroundStyle(ZyquoTheme.textTertiary)
            }
        }
        .frame(maxWidth: .infinity)
    }

    private func column(side: Int) -> some View {
        ScrollView {
            VStack(alignment: .leading, spacing: ZyquoTheme.Spacing.s) {
                if !controller.thinking[side].isEmpty {
                    ThinkingDisclosure(thinking: controller.thinking[side], isLive: false)
                }
                MarkdownText(controller.outputs[side])
                    .textSelection(.enabled)
                if controller.generating[side] {
                    ProgressView().controlSize(.small)
                }
            }
            .padding(ZyquoTheme.Spacing.m)
            .frame(maxWidth: .infinity, alignment: .leading)
        }
        .frame(maxWidth: .infinity)
    }

    private var inputRow: some View {
        HStack(spacing: ZyquoTheme.Spacing.xs) {
            TextField("Prompt both models…", text: $prompt, axis: .vertical)
                .textFieldStyle(.plain)
                .font(ZyquoTheme.chatBody)
                .lineLimit(1...6)
                .focused($focused)
                .onSubmit(broadcast)
            Button {
                broadcast()
            } label: {
                Image(systemName: "arrow.up")
                    .font(.system(size: 13, weight: .bold))
                    .foregroundStyle(.white)
                    .frame(width: 28, height: 28)
                    .background(canSend ? ZyquoTheme.accent : ZyquoTheme.textTertiary, in: Circle())
            }
            .buttonStyle(PressableButtonStyle())
            .keyboardShortcut(.return, modifiers: .command)
            .disabled(!canSend)
        }
        .padding(ZyquoTheme.Spacing.s)
        .background(ZyquoTheme.surface, in: RoundedRectangle(cornerRadius: ZyquoTheme.Radius.l))
        .overlay(
            RoundedRectangle(cornerRadius: ZyquoTheme.Radius.l)
                .stroke(ZyquoTheme.border, lineWidth: ZyquoTheme.hairline)
        )
        .floatingShadow()
        .padding(ZyquoTheme.Spacing.m)
    }

    private var canSend: Bool {
        !prompt.trimmingCharacters(in: .whitespaces).isEmpty
            && controller.states.allSatisfy { if case .ready = $0 { true } else { false } }
            && !controller.generating.contains(true)
    }

    private func broadcast() {
        guard canSend else { return }
        let text = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
        prompt = ""
        controller.broadcast(prompt: text, params: app.settings.defaultParams)
    }
}

/// Owns two independent engines for compare mode. The main app engine is
/// left untouched; RAM gating warns when both models exceed the safe budget.
@MainActor
@Observable
final class CompareController {
    enum ColumnState {
        case empty
        case loading
        case ready
    }

    var selection: [String?] = [nil, nil]
    var states: [ColumnState] = [.empty, .empty]
    var outputs: [String] = ["", ""]
    var thinking: [String] = ["", ""]
    var generating: [Bool] = [false, false]
    var stats: [GenerationStats?] = [nil, nil]
    var ramWarning: String?

    private var engines: [InferenceEngine?] = [nil, nil]
    private var sessions: [Conversation] = [Conversation(), Conversation()]

    func loadSelection(app: AppModel) async {
        // RAM gate across both columns.
        let sizes = selection.compactMap { id in id.flatMap { app.store.model(for: $0)?.sizeBytes } }
        let combined = sizes.reduce(0, +)
        switch MemoryAdvisor.verdict(weightsBytes: combined) {
        case .fits:
            ramWarning = nil
        case .tight:
            ramWarning = "Both models together are a tight fit for this Mac's memory."
        case .tooLarge:
            ramWarning = "These two models do not fit in memory together — pick smaller ones."
            return
        }

        for side in 0..<2 {
            guard let repoID = selection[side], let model = app.store.model(for: repoID) else {
                engines[side] = nil
                states[side] = .empty
                continue
            }
            if let engine = engines[side], await engine.currentModel?.repoID == repoID { continue }
            states[side] = .loading
            let engine = InferenceEngine()
            do {
                try await engine.load(model: model)
                sessions[side] = Conversation(modelID: repoID)
                try await engine.startSession(conversation: sessions[side])
                engines[side] = engine
                states[side] = .ready
            } catch {
                states[side] = .empty
                ramWarning = error.localizedDescription
            }
        }
    }

    func broadcast(prompt: String, params: GenerationParams) {
        for side in 0..<2 {
            guard let engine = engines[side] else { continue }
            outputs[side] = ""
            thinking[side] = ""
            stats[side] = nil
            generating[side] = true
            Task {
                var parser = ThinkTagParser()
                do {
                    let events = try await engine.generate(prompt: prompt, params: params)
                    for try await event in events {
                        switch event {
                        case .token(let t):
                            let (visible, think, _) = parser.consume(t)
                            if !visible.isEmpty { outputs[side] += visible }
                            if !think.isEmpty { thinking[side] += think }
                        case .stats(let s):
                            stats[side] = s
                        case .finished:
                            break
                        }
                    }
                } catch {
                    outputs[side] += "\n*\(error.localizedDescription)*"
                }
                generating[side] = false
            }
        }
    }

    func teardown() async {
        for engine in engines.compactMap({ $0 }) {
            await engine.unload()
        }
        engines = [nil, nil]
        states = [.empty, .empty]
    }
}