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

import MLXLMCommon
import SwiftUI
import UniformTypeIdentifiers

/// Playground: interactive inference per model type — streaming chat for
/// LLM/VLM (images attach for VLM), a text→vector inspector with similarity
/// for embeddings, per-run stats, and verifiable load/unload.
struct PlaygroundView: View {
    @Bindable var model: AppModel
    @State private var session = PlaygroundSession()

    var body: some View {
        if model.models.isEmpty {
            EmptyStateView(
                icon: "bubble.left.and.text.bubble.right",
                title: "The Playground Awaits",
                message: "Load any local model and interact with it — streaming chat for LLMs, images for vision models, a vector inspector for embeddings — with live tokens/sec, time-to-first-token, and memory stats.")
        } else {
            VStack(spacing: 0) {
                playgroundHeader
                Rectangle().fill(ZyquoTheme.border).frame(height: ZyquoTheme.hairline)
                content
            }
        }
    }

    private var playgroundHeader: some View {
        HStack(spacing: ZyquoTheme.spacing12) {
            Picker("Model", selection: $session.selectedModelID) {
                Text("Choose a model…").tag(String?.none)
                ForEach(model.models.filter { $0.type.isSwiftNative || $0.type == .speech }) { item in
                    Text("\(item.name)  (\(item.type.displayName))").tag(String?.some(item.id))
                }
            }
            .frame(maxWidth: 420)

            if session.isLoading {
                ProgressView().controlSize(.small)
                Text("Loading…")
                    .font(ZyquoTheme.captionFont)
                    .foregroundStyle(ZyquoTheme.textSecondary)
            } else if session.loadedModel != nil {
                StatusPill(text: "Loaded", color: ZyquoTheme.success)
                Button("Unload") {
                    Task { await session.unload() }
                }
                .controlSize(.small)
                if let freed = session.lastFreedBytes {
                    Text("freed \(ByteCountFormatter.string(fromByteCount: freed, countStyle: .memory))")
                        .font(ZyquoTheme.captionFont)
                        .foregroundStyle(ZyquoTheme.textTertiary)
                }
            }

            Spacer()

            if let stats = session.lastStats {
                Text(String(format: "%.1f tok/s · TTFT %.2fs · %d tokens",
                    stats.tokensPerSecond, stats.ttft, stats.generatedTokens))
                    .font(ZyquoTheme.monoSmallFont)
                    .foregroundStyle(ZyquoTheme.textSecondary)
            }
        }
        .padding(.horizontal, ZyquoTheme.spacing20)
        .padding(.vertical, ZyquoTheme.spacing8)
        .onChange(of: session.selectedModelID) { _, newID in
            guard let newID, let item = model.models.first(where: { $0.id == newID }) else { return }
            if item.type == .speech {
                // Out-of-process — nothing to load in the engine.
                session.speechModel = item
            } else {
                Task { await session.load(item) }
            }
        }
    }

    @ViewBuilder
    private var content: some View {
        switch session.loadedModel?.type ?? selectedType {
        case .embedding:
            EmbeddingInspector(session: session)
        case .llm, .vlm:
            ChatPanel(session: session)
        case .speech:
            SpeechPanel(session: session)
        default:
            EmptyStateView(
                icon: "cpu",
                title: "Pick a Model to Begin",
                message: "Choose an installed model above. LLMs and vision models open a streaming chat; embedding models open the vector inspector; speech models open the transcriber.")
        }
    }

    /// Speech models run out-of-process (no in-engine load) — panel selection
    /// falls back to the picker choice.
    private var selectedType: ModelType? {
        model.models.first { $0.id == session.selectedModelID }?.type
    }
}

// MARK: - Session state

@Observable
@MainActor
final class PlaygroundSession {
    var selectedModelID: String?
    var loadedModel: LocalModel?
    var isLoading = false
    var isGenerating = false
    var lastStats: InferenceStats?
    var lastFreedBytes: Int64?
    var errorMessage: String?

    // Chat state
    var transcript: [(role: String, text: String)] = []
    var prompt = ""
    var attachedImage: URL?
    private var generationTask: Task<Void, Never>?

    // Embedding state
    var embedInput = "The quick brown fox\nA fast auburn fox\nQuarterly revenue grew 4%"
    var embedResults: [(text: String, vector: [Float])] = []
    var similarities: [(a: String, b: String, score: Float)] = []

    // Speech state (out-of-process — no engine load)
    var speechModel: LocalModel?
    var transcription: TranscriptionResult?
    var isTranscribing = false

    func transcribe(audio: URL) {
        guard let speechModel else { return }
        isTranscribing = true
        transcription = nil
        Task {
            do {
                transcription = try await SpeechService.shared.transcribe(
                    model: speechModel, audio: audio)
            } catch {
                errorMessage = error.localizedDescription
            }
            isTranscribing = false
        }
    }

    func load(_ model: LocalModel) async {
        isLoading = true
        errorMessage = nil
        transcript = []
        embedResults = []
        similarities = []
        lastStats = nil
        do {
            try await InferenceEngine.shared.load(model: model)
            loadedModel = model
        } catch {
            errorMessage = error.localizedDescription
            loadedModel = nil
        }
        isLoading = false
    }

    func unload() async {
        stop()
        lastFreedBytes = try? await InferenceEngine.shared.unload()
        loadedModel = nil
        selectedModelID = nil
    }

    func send() {
        let text = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !text.isEmpty, !isGenerating, loadedModel != nil else { return }
        prompt = ""
        let image = attachedImage
        attachedImage = nil
        transcript.append((role: "user", text: text))
        transcript.append((role: "assistant", text: ""))
        isGenerating = true

        var messages: [Chat.Message] = transcript.dropLast(2).map {
            $0.role == "user" ? .user($0.text) : .assistant($0.text)
        }
        if let image {
            messages.append(.user(text, images: [.url(image)]))
        } else {
            messages.append(.user(text))
        }

        generationTask = Task {
            do {
                let stream = try await InferenceEngine.shared.generate(
                    messages: messages, params: GenerationParams())
                for try await event in stream {
                    switch event {
                    case .chunk(let piece):
                        transcript[transcript.count - 1].text += piece
                    case .finished(let stats):
                        lastStats = stats
                    }
                }
            } catch {
                errorMessage = error.localizedDescription
            }
            isGenerating = false
        }
    }

    func stop() {
        generationTask?.cancel()
        generationTask = nil
        isGenerating = false
    }

    func runEmbedding() {
        let texts = embedInput.split(separator: "\n").map(String.init).filter { !$0.isEmpty }
        guard !texts.isEmpty else { return }
        Task {
            do {
                let vectors = try await InferenceEngine.shared.embed(texts: texts)
                embedResults = Array(zip(texts, vectors))
                similarities = []
                for i in 0..<vectors.count {
                    for j in (i + 1)..<vectors.count {
                        similarities.append((
                            a: texts[i], b: texts[j],
                            score: InferenceEngine.cosineSimilarity(vectors[i], vectors[j])))
                    }
                }
                similarities.sort { $0.score > $1.score }
            } catch {
                errorMessage = error.localizedDescription
            }
        }
    }
}

// MARK: - Chat panel

private struct ChatPanel: View {
    @Bindable var session: PlaygroundSession
    @State private var isPickingImage = false

    var body: some View {
        VStack(spacing: 0) {
            if let error = session.errorMessage {
                Text(error)
                    .font(ZyquoTheme.captionFont)
                    .foregroundStyle(ZyquoTheme.danger)
                    .padding(ZyquoTheme.spacing8)
            }

            ScrollViewReader { proxy in
                ScrollView {
                    LazyVStack(alignment: .leading, spacing: ZyquoTheme.spacing12) {
                        ForEach(Array(session.transcript.enumerated()), id: \.offset) { index, entry in
                            MessageBubble(role: entry.role, text: entry.text)
                                .id(index)
                        }
                    }
                    .padding(ZyquoTheme.spacing20)
                }
                .onChange(of: session.transcript.last?.text) {
                    proxy.scrollTo(session.transcript.count - 1, anchor: .bottom)
                }
            }

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

            HStack(spacing: ZyquoTheme.spacing8) {
                if session.loadedModel?.type == .vlm {
                    Button {
                        isPickingImage = true
                    } label: {
                        Image(systemName: session.attachedImage == nil ? "photo" : "photo.fill")
                            .foregroundStyle(
                                session.attachedImage == nil
                                    ? ZyquoTheme.textSecondary : ZyquoTheme.accent)
                    }
                    .buttonStyle(.plain)
                    .help(session.attachedImage?.lastPathComponent ?? "Attach an image")
                }

                TextField("Message the model…", text: $session.prompt, axis: .vertical)
                    .textFieldStyle(.plain)
                    .font(ZyquoTheme.bodyFont)
                    .lineLimit(1...5)
                    .onSubmit { session.send() }

                if session.isGenerating {
                    Button {
                        session.stop()
                    } label: {
                        Image(systemName: "stop.circle.fill")
                            .font(.system(size: 20))
                            .foregroundStyle(ZyquoTheme.danger)
                    }
                    .buttonStyle(.plain)
                } else {
                    Button {
                        session.send()
                    } label: {
                        Image(systemName: "arrow.up.circle.fill")
                            .font(.system(size: 20))
                            .foregroundStyle(
                                session.prompt.isEmpty ? ZyquoTheme.textTertiary : ZyquoTheme.accent)
                    }
                    .buttonStyle(.plain)
                    .disabled(session.prompt.isEmpty)
                }
            }
            .padding(ZyquoTheme.spacing12)
            .background(ZyquoTheme.surface)
        }
        .fileImporter(isPresented: $isPickingImage, allowedContentTypes: [.image]) { result in
            if case .success(let url) = result { session.attachedImage = url }
        }
    }
}

private struct MessageBubble: View {
    let role: String
    let text: String

    var body: some View {
        HStack {
            if role == "user" { Spacer(minLength: 80) }
            Text(text.isEmpty ? "…" : text)
                .font(ZyquoTheme.bodyFont)
                .foregroundStyle(ZyquoTheme.textPrimary)
                .textSelection(.enabled)
                .padding(ZyquoTheme.spacing12)
                .background(role == "user" ? ZyquoTheme.accentSubtle : ZyquoTheme.surface)
                .clipShape(RoundedRectangle(cornerRadius: ZyquoTheme.radiusMedium))
                .overlay(
                    RoundedRectangle(cornerRadius: ZyquoTheme.radiusMedium)
                        .stroke(ZyquoTheme.border, lineWidth: ZyquoTheme.hairline))
            if role != "user" { Spacer(minLength: 80) }
        }
    }
}

// MARK: - Speech panel

private struct SpeechPanel: View {
    @Bindable var session: PlaygroundSession
    @State private var isPickingAudio = false

    var body: some View {
        VStack(spacing: ZyquoTheme.spacing16) {
            if session.isTranscribing {
                ProgressView("Transcribing…")
            } else if let result = session.transcription {
                ScrollView {
                    VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {
                        HStack {
                            if let language = result.language {
                                StatusPill(text: language, color: ZyquoTheme.slate)
                            }
                            Text("\(result.segments) segments · \(String(format: "%.1fs", result.duration))")
                                .font(ZyquoTheme.captionFont)
                                .foregroundStyle(ZyquoTheme.textSecondary)
                        }
                        Text(result.text)
                            .font(ZyquoTheme.bodyFont)
                            .textSelection(.enabled)
                    }
                    .padding(ZyquoTheme.spacing16)
                    .frame(maxWidth: .infinity, alignment: .leading)
                    .zyquoCard()
                    .padding(ZyquoTheme.spacing20)
                }
            } else {
                EmptyStateView(
                    icon: "waveform",
                    title: "Transcribe Audio",
                    message: "Pick an audio file — the speech model runs locally through the Python pipeline and returns the transcript here.",
                    actionLabel: "Choose Audio…",
                    action: { isPickingAudio = true })
            }

            if session.transcription != nil {
                Button("Transcribe Another…") { isPickingAudio = true }
                    .padding(.bottom, ZyquoTheme.spacing16)
            }
        }
        .fileImporter(isPresented: $isPickingAudio, allowedContentTypes: [.audio]) { result in
            if case .success(let url) = result { session.transcribe(audio: url) }
        }
    }
}

// MARK: - Embedding inspector

private struct EmbeddingInspector: View {
    @Bindable var session: PlaygroundSession

    var body: some View {
        ScrollView {
            VStack(alignment: .leading, spacing: ZyquoTheme.spacing16) {
                Text("One text per line — embed them and compare similarities.")
                    .font(ZyquoTheme.captionFont)
                    .foregroundStyle(ZyquoTheme.textSecondary)

                TextEditor(text: $session.embedInput)
                    .font(ZyquoTheme.monoFont)
                    .frame(height: 120)
                    .padding(ZyquoTheme.spacing8)
                    .background(ZyquoTheme.surfaceSecondary)
                    .clipShape(RoundedRectangle(cornerRadius: ZyquoTheme.radiusSmall))

                Button("Embed") { session.runEmbedding() }
                    .buttonStyle(.borderedProminent)
                    .tint(ZyquoTheme.accent)

                if !session.embedResults.isEmpty {
                    VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {
                        Text("Vectors")
                            .font(ZyquoTheme.headlineFont)
                        ForEach(session.embedResults, id: \.text) { result in
                            HStack {
                                Text("dim \(result.vector.count)")
                                    .font(ZyquoTheme.monoSmallFont)
                                    .foregroundStyle(ZyquoTheme.chartThroughput)
                                Text(result.vector.prefix(6)
                                    .map { String(format: "%+.3f", $0) }
                                    .joined(separator: " ") + " …")
                                    .font(ZyquoTheme.monoSmallFont)
                                    .foregroundStyle(ZyquoTheme.textSecondary)
                                Text(result.text)
                                    .font(ZyquoTheme.captionFont)
                                    .foregroundStyle(ZyquoTheme.textPrimary)
                                    .lineLimit(1)
                            }
                        }
                    }
                    .padding(ZyquoTheme.spacing12)
                    .frame(maxWidth: .infinity, alignment: .leading)
                    .zyquoCard()

                    VStack(alignment: .leading, spacing: ZyquoTheme.spacing8) {
                        Text("Cosine similarity")
                            .font(ZyquoTheme.headlineFont)
                        ForEach(Array(session.similarities.enumerated()), id: \.offset) { _, pair in
                            HStack {
                                Text(String(format: "%.4f", pair.score))
                                    .font(ZyquoTheme.monoFont)
                                    .foregroundStyle(
                                        pair.score > 0.8 ? ZyquoTheme.success : ZyquoTheme.textSecondary)
                                Text("\(pair.a)  ↔  \(pair.b)")
                                    .font(ZyquoTheme.captionFont)
                                    .foregroundStyle(ZyquoTheme.textPrimary)
                                    .lineLimit(1)
                            }
                        }
                    }
                    .padding(ZyquoTheme.spacing12)
                    .frame(maxWidth: .infinity, alignment: .leading)
                    .zyquoCard()
                }
            }
            .padding(ZyquoTheme.spacing20)
        }
    }
}