//
//  RunController.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Owns the live run for one task: assembles the engine (provider client +
//  key, workspace, tools, policy gate with the UI approval presenter, audit
//  log, AgentLoop) and consumes the AsyncThrowingStream<AgentEvent>,
//  publishing everything the command-center UI renders — live step cards
//  (streaming text/thinking/tool arguments/tool output), the plan, budgets,
//  the pending approval, guard trips, compactions, the terminal feed, audit
//  entries, and workspace files. AgentEvents arrive off the main thread; the
//  consuming Task is MainActor-bound so every mutation happens on main.
//

import Combine
import Foundation

// MARK: - Live timeline model

/// One tool invocation as it renders live inside a step card.
struct LiveInvocation: Identifiable {
    enum Phase {
        /// Arguments still streaming from the model.
        case streaming
        /// Cleared (or clearing) the policy gate and executing.
        case executing
        case finished
    }

    var id: String
    var streamIndex: Int?
    var name: String
    /// Accumulating raw argument JSON (may be partial while streaming).
    var argumentsJSON: String
    var phase: Phase = .streaming
    var outputLines: [TerminalLine] = []
    var result: ToolResult?
    var exitCode: Int32?
    var policyDecision: PolicyDecisionRecord?

    /// The human-facing payload: the shell command / script / path when the
    /// arguments parse, else the raw JSON.
    var displayPayload: String {
        guard let data = argumentsJSON.data(using: .utf8),
              let object = try? JSONSerialization.jsonObject(with: data) as? [String: Any] else {
            return argumentsJSON
        }
        for key in ["command", "script", "path"] {
            if let value = object[key] as? String { return value }
        }
        return argumentsJSON
    }
}

/// One agent step as it renders live (mirrors AgentStep, mutable per event).
struct LiveStep: Identifiable {
    var id: UUID
    var index: Int
    var status: AgentStepStatus
    var thinking: String
    var text: String
    var invocations: [LiveInvocation]
    var startedAt: Date
    var finishedAt: Date?
    var inputTokens: Int?
    var outputTokens: Int?

    /// A step with no tool calls whose run completed = the final answer.
    var isFinalAnswer: Bool { invocations.isEmpty && status == .completed && !text.isEmpty }

    init(step: AgentStep) {
        self.id = step.id
        self.index = step.index
        self.status = step.status
        self.thinking = step.thinking ?? ""
        self.text = step.text
        self.invocations = step.toolInvocations.map { LiveInvocation(invocation: $0) }
        self.startedAt = step.startedAt
        self.finishedAt = step.finishedAt
        self.inputTokens = step.inputTokens
        self.outputTokens = step.outputTokens
    }
}

extension LiveInvocation {
    /// Builds a finished invocation view from a persisted AgentToolInvocation
    /// (past-run re-rendering and stepCompleted reconciliation).
    init(invocation: AgentToolInvocation) {
        self.id = invocation.call.id
        self.streamIndex = nil
        self.name = invocation.call.name
        self.argumentsJSON = invocation.call.argumentsJSON
        self.phase = .finished
        self.result = invocation.result
        self.exitCode = invocation.exitCode
        self.policyDecision = invocation.policyDecision
    }
}

/// One item of the live run timeline (step cards interleaved with
/// compaction notices, in event order).
enum RunEntry: Identifiable {
    case step(LiveStep)
    case compaction(CompactionRecord)

    var id: UUID {
        switch self {
        case .step(let step): return step.id
        case .compaction(let record): return record.id
        }
    }
}

/// One line of the Activity/Terminal drawer's live feed.
struct TerminalLine: Identifiable {
    enum Kind {
        case command
        case stdout
        case stderr
        case note
        /// cwd banner / run lifecycle marker.
        case meta
    }

    let id = UUID()
    var kind: Kind
    var text: String
    var timestamp: Date = Date()
}

// MARK: - RunController

@MainActor
final class RunController: ObservableObject {
    // Live run state
    @Published private(set) var entries: [RunEntry] = []
    @Published private(set) var plan: TaskPlan?
    @Published private(set) var runStatus: AgentRunStatus?
    @Published private(set) var outcome: AgentRunOutcome?
    @Published private(set) var pendingApproval: PendingApproval?
    @Published private(set) var guardTrip: LoopGuardTrip?
    @Published private(set) var isRunning = false
    @Published private(set) var lastError: String?

    // Budgets & usage (Plan panel meters)
    @Published private(set) var tokensUsed = 0
    @Published private(set) var stepsUsed = 0
    @Published private(set) var runStartedAt: Date?
    @Published private(set) var loopGuardConfiguration = LoopGuardConfiguration.default

    // Drawer feeds
    @Published private(set) var terminalLines: [TerminalLine] = []
    @Published private(set) var auditEntries: [AuditEntry] = []
    @Published private(set) var workspaceFiles: [WorkspaceFileEntry] = []

    // Info popover
    @Published private(set) var systemPromptPreview: String?

    let taskID: AgentTask.ID
    private let store: TaskStore
    private let policyPersistence: PersistenceService
    /// App-wide agent settings (budgets, timeouts, persona-independent
    /// tunables). Nil in headless smoke tests → AgentConfiguration.default.
    private weak var settings: AgentSettingsStore?

    private var loop: AgentLoop?
    private var policy: PolicyEngine?
    private var audit: AuditLog?
    private var workspace: WorkspaceManager?
    private var consumeTask: Task<Void, Never>?
    /// Retained for the post-run auto-title call (app runs only).
    private var titleContext: (model: AIModel, client: any ProviderClient, apiKey: String)?
    private var titleTask: Task<Void, Never>?

    /// Terminal feed cap — append-only ring so hours-long runs stay light.
    private static let terminalLineCap = 4000

    init(
        taskID: AgentTask.ID,
        store: TaskStore,
        policyPersistence: PersistenceService = .shared,
        settings: AgentSettingsStore? = nil
    ) {
        self.taskID = taskID
        self.store = store
        self.policyPersistence = policyPersistence
        self.settings = settings
        if let workspaceURL = store.task(id: taskID)?.workspaceURL,
           let attached = try? WorkspaceManager(existingAt: workspaceURL) {
            self.workspace = attached
            self.audit = AuditLog(fileURL: attached.internalDirectory.appendingPathComponent("audit.jsonl"))
            self.plan = store.task(id: taskID)?.messages.last(where: { $0.plan != nil })?.plan
            refreshWorkspaceState()
        }
    }

    private var task: AgentTask? { store.task(id: taskID) }

    // MARK: - Starting a run

    /// Starts a run for the given prompt. `client`/`apiKey` are injectable
    /// for the offline UI smoke test; app runs resolve them from the
    /// ProviderRegistry and the encrypted vault / environment.
    func start(
        prompt: String,
        model: AIModel,
        persona: Persona? = nil,
        client injectedClient: (any ProviderClient)? = nil,
        apiKey injectedKey: String? = nil
    ) {
        guard !isRunning, var task else { return }
        let trimmedPrompt = prompt.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmedPrompt.isEmpty else { return }

        guard model.capabilities.tools else {
            lastError = "\(model.displayName) does not support tool calling and cannot run agent tasks. Pick an agent-capable model."
            return
        }

        let client = injectedClient ?? ProviderRegistry.client(for: model)
        let apiKey: String
        if let injectedKey {
            apiKey = injectedKey
        } else if let resolved = AgentCLI.resolveAPIKey(for: model.provider) {
            apiKey = resolved
        } else {
            lastError = ProviderError.missingAPIKey(model.provider).localizedDescription
            return
        }

        // Workspace: reattach the task's existing one, or create it now.
        let workspace: WorkspaceManager
        do {
            if let url = task.workspaceURL {
                workspace = try WorkspaceManager(existingAt: url)
            } else {
                let title = task.title == "New Task" ? AgentTask.title(fromPrompt: trimmedPrompt) : task.title
                workspace = try WorkspaceManager(taskTitle: title)
                task.workspacePath = workspace.root.path
            }
        } catch {
            lastError = "Could not prepare the workspace: \(error.localizedDescription)"
            return
        }
        self.workspace = workspace

        // First prompt names the task.
        if task.title == "New Task" {
            task.title = AgentTask.title(fromPrompt: trimmedPrompt)
        }
        task.modelID = model.id
        task.providerID = model.provider
        task.messages.append(TaskMessage(kind: .user, text: trimmedPrompt))
        task.status = .planning
        store.update(task)

        // Engine assembly — every action passes the policy gate + audit log.
        let presenter = UIApprovalPresenter()
        presenter.onRequest = { [weak self] pending in
            Task { @MainActor in
                self?.presentApproval(pending)
            }
        }
        let policy = PolicyEngine(mode: task.safetyMode, approvals: presenter, persistence: policyPersistence)
        let audit = AuditLog(fileURL: workspace.internalDirectory.appendingPathComponent("audit.jsonl"))

        // Settings › Agent budgets/tunables + the persona's prompt addendum.
        let personaAddendum = persona?.systemPrompt
        let configuration = settings?.agentConfiguration(personaAddendum: personaAddendum)
            ?? { var c = AgentConfiguration.default; c.personaAddendum = personaAddendum; return c }()
        var executionConfiguration = ExecutionConfiguration.default
        if let timeout = configuration.perCommandTimeout {
            executionConfiguration.defaultTimeout = timeout
        }
        self.loopGuardConfiguration = configuration.loopGuard

        let executor = ExecutionService(configuration: executionConfiguration)
        let tools = ToolRegistry.standard(executor: executor)
        let loop = AgentLoop(
            model: model,
            client: client,
            apiKey: apiKey,
            tools: tools,
            policy: policy,
            audit: audit,
            workspace: workspace,
            configuration: configuration
        )
        self.policy = policy
        self.audit = audit
        self.loop = loop
        // Auto-title uses the injected client only for real runs — smoke
        // tests inject a scripted client and must stay deterministic.
        self.titleContext = injectedClient == nil ? (model, client, apiKey) : nil
        self.systemPromptPreview = AgentSystemPrompt.build(
            workspacePath: workspace.root.path,
            toolNames: tools.toolNames,
            safetyMode: task.safetyMode,
            personaAddendum: personaAddendum
        )

        // Reset live state.
        entries = []
        outcome = nil
        guardTrip = nil
        pendingApproval = nil
        lastError = nil
        tokensUsed = 0
        stepsUsed = 0
        isRunning = true
        runStartedAt = Date()
        appendTerminal(.meta, "▶ run started — \(model.displayName) · \(task.safetyMode.displayName) mode")
        appendTerminal(.meta, "cwd \(workspace.root.path)")

        // Consume the event stream on the MainActor (the stream itself is
        // produced inside the AgentLoop actor; only the handling hops here).
        consumeTask = Task { [weak self] in
            do {
                for try await event in await loop.run(task: trimmedPrompt) {
                    self?.handle(event)
                }
            } catch {
                self?.finish(with: .failed(reason: error.localizedDescription))
            }
        }
    }

    // MARK: - Run control

    /// Stops the run (⌘. / Stop button): kills the in-flight model stream
    /// and any executing process.
    func cancel() {
        pendingApproval?.resolve(.deny)
        pendingApproval = nil
        let loop = loop
        Task { await loop?.cancel() }
    }

    /// Answers a guard trip with "continue" (optionally raising budgets).
    func resumeAfterTrip(raisingBudget: Bool = true) {
        guardTrip = nil
        updateTaskStatus(.running)
        let loop = loop
        Task { await loop?.resume(raisingBudget: raisingBudget) }
    }

    /// Answers a guard trip with "stop".
    func stopAfterTrip() {
        guardTrip = nil
        let loop = loop
        Task { await loop?.stop() }
    }

    /// Resolves the pending approval card.
    func resolveApproval(_ resolution: ApprovalResolution) {
        guard let pending = pendingApproval else { return }
        pendingApproval = nil
        mutateLastStep { step in
            if step.status == .awaitingApproval { step.status = .executing }
        }
        if isRunning {
            updateTaskStatus(.running)
        }
        pending.resolve(resolution)
    }

    /// Switches the safety mode for this task (live runs switch immediately).
    func setSafetyMode(_ mode: SafetyMode) {
        guard var task, task.safetyMode != mode else { return }
        task.safetyMode = mode
        store.update(task, touch: false)
        let policy = policy
        Task { await policy?.setMode(mode) }
    }

    // MARK: - Event handling (MainActor)

    private func handle(_ event: AgentEvent) {
        switch event {
        case .statusChanged(let status):
            runStatus = status
            switch status {
            case .preparing:
                updateTaskStatus(.planning)
            case .running, .executingTools, .compacting:
                if pendingApproval == nil { updateTaskStatus(.running) }
            case .awaitingUser:
                updateTaskStatus(.awaitingInput)
            case .finished:
                break // runFinished carries the outcome
            }

        case .stepStarted(let step):
            entries.append(.step(LiveStep(step: step)))
            stepsUsed = max(stepsUsed, step.index)

        case .thinkingDelta(let stepID, let delta):
            mutateStep(id: stepID) { $0.thinking += delta }

        case .textDelta(let stepID, let delta):
            mutateStep(id: stepID) { $0.text += delta }

        case .toolCallStreaming(let stepID, let index, let id, let name):
            mutateStep(id: stepID) { step in
                var invocation = LiveInvocation(id: id, streamIndex: index, name: name, argumentsJSON: "")
                invocation.phase = .streaming
                step.invocations.append(invocation)
            }

        case .toolCallArgumentsDelta(let stepID, let index, let delta):
            mutateStep(id: stepID) { step in
                if let i = step.invocations.lastIndex(where: { $0.streamIndex == index }) {
                    step.invocations[i].argumentsJSON += delta
                }
            }

        case .toolCallStarted(let stepID, let invocation):
            mutateStep(id: stepID) { step in
                if let i = step.invocations.firstIndex(where: { $0.id == invocation.call.id }) {
                    step.invocations[i].argumentsJSON = invocation.call.argumentsJSON
                    step.invocations[i].phase = .executing
                } else {
                    var live = LiveInvocation(
                        id: invocation.call.id,
                        streamIndex: nil,
                        name: invocation.call.name,
                        argumentsJSON: invocation.call.argumentsJSON
                    )
                    live.phase = .executing
                    step.invocations.append(live)
                }
                step.status = .executing
            }
            if let payload = payloadPreview(of: invocation.call) {
                appendTerminal(.command, "$ \(payload)")
            }

        case .toolOutput(let stepID, let callID, let chunk):
            let line: TerminalLine
            switch chunk {
            case .stdout(let text): line = TerminalLine(kind: .stdout, text: text)
            case .stderr(let text): line = TerminalLine(kind: .stderr, text: text)
            case .note(let text): line = TerminalLine(kind: .note, text: text)
            }
            mutateStep(id: stepID) { step in
                if let i = step.invocations.firstIndex(where: { $0.id == callID }) {
                    step.invocations[i].outputLines.append(line)
                }
            }
            appendTerminal(line.kind, line.text)

        case .toolCallFinished(let stepID, let invocation):
            mutateStep(id: stepID) { step in
                if let i = step.invocations.firstIndex(where: { $0.id == invocation.call.id }) {
                    step.invocations[i].phase = .finished
                    step.invocations[i].result = invocation.result
                    step.invocations[i].exitCode = invocation.exitCode
                    step.invocations[i].policyDecision = invocation.policyDecision
                } else {
                    step.invocations.append(LiveInvocation(invocation: invocation))
                }
            }

        case .stepCompleted(let step):
            mutateStep(id: step.id) { live in
                live.status = step.status
                live.finishedAt = step.finishedAt
                live.inputTokens = step.inputTokens
                live.outputTokens = step.outputTokens
                if !step.text.isEmpty { live.text = step.text }
                if let thinking = step.thinking { live.thinking = thinking }
                for invocation in step.toolInvocations {
                    if let i = live.invocations.firstIndex(where: { $0.id == invocation.call.id }) {
                        live.invocations[i].result = invocation.result
                        live.invocations[i].exitCode = invocation.exitCode
                        live.invocations[i].policyDecision = invocation.policyDecision
                        live.invocations[i].phase = .finished
                    }
                }
            }
            tokensUsed += (step.inputTokens ?? 0) + (step.outputTokens ?? 0)
            refreshWorkspaceState()

        case .planUpdated(let updated):
            plan = updated

        case .guardTripped(let trip):
            guardTrip = trip
            appendTerminal(.meta, "⏸ loop guard [\(trip.reason.rawValue)] — \(trip.message)")

        case .compactionPerformed(let record):
            entries.append(.compaction(record))
            appendTerminal(.meta, "⟳ compacted \(record.summarizedSteps) step(s): ~\(record.beforeTokens) → ~\(record.afterTokens) tokens")

        case .runFinished(let outcome):
            finish(with: outcome)
        }
    }

    private func presentApproval(_ pending: PendingApproval) {
        pendingApproval = pending
        mutateLastStep { $0.status = .awaitingApproval }
        updateTaskStatus(.awaitingApproval)
        appendTerminal(.meta, "⚠ approval required: \(pending.action.payload)")
    }

    private func finish(with outcome: AgentRunOutcome) {
        self.outcome = outcome
        isRunning = false
        consumeTask = nil
        loop = nil

        // Fold the finished run into the task's persisted history so past
        // runs re-render exactly, then clear the live timeline.
        var steps: [AgentStep] = []
        if let workspaceURL = task?.workspaceURL,
           let transcript = Transcript.load(from: workspaceURL) {
            steps = transcript.steps
        }

        let finalText: String
        let status: AgentTaskStatus
        switch outcome {
        case .completed(let answer):
            finalText = answer
            status = .done
            appendTerminal(.meta, "✔ run completed")
        case .failed(let reason):
            finalText = reason
            status = .failed
            appendTerminal(.meta, "✘ run failed — \(reason)")
        case .cancelled:
            finalText = "Run cancelled."
            status = .idle
            appendTerminal(.meta, "■ run cancelled")
        case .stoppedByUser(let reason):
            finalText = "Run stopped — \(reason)"
            status = .idle
            appendTerminal(.meta, "■ run stopped — \(reason)")
        }

        if var task {
            task.messages.append(TaskMessage(
                kind: .agentRun,
                text: finalText,
                steps: steps,
                plan: plan,
                outcome: outcome
            ))
            task.status = status
            store.update(task)
        }
        entries = []
        pendingApproval = nil
        guardTrip = nil
        refreshWorkspaceState()
        refreshAudit()

        if case .completed = outcome {
            generateTitleIfStillDefault()
        }
    }

    // MARK: - Auto-title

    /// After a successful run, if the title is still the prompt-derived
    /// default, fire a cheap background call (the task's own provider/model,
    /// ≤20 tokens) to produce a 4–6 word title. Silent failure by design.
    private func generateTitleIfStillDefault() {
        guard titleTask == nil, let context = titleContext, let task else { return }
        guard let firstPrompt = task.messages.first(where: { $0.kind == .user })?.text,
              task.title == AgentTask.title(fromPrompt: firstPrompt) else { return }
        let request = ChatRequest(
            model: context.model,
            systemPrompt: "You write concise task titles.",
            messages: [Message(
                role: .user,
                text: "Summarize this task as a 4-6 word title. Reply with the title only — no quotes, no trailing punctuation.\n\nTask: \(firstPrompt.prefix(600))"
            )],
            parameters: ChatParameters(maxTokens: 20),
            stream: false
        )
        let taskID = taskID
        let store = store
        titleTask = Task { [weak self] in
            defer { self?.titleTask = nil }
            guard let reply = try? await context.client.complete(request, apiKey: context.apiKey) else { return }
            let title = reply.text
                .trimmingCharacters(in: .whitespacesAndNewlines)
                .trimmingCharacters(in: CharacterSet(charactersIn: "\"'“”.\n"))
            guard !title.isEmpty, title.count <= 80, !title.contains("\n") else { return }
            store.rename(taskID, to: title)
        }
    }

    // MARK: - Drawer refresh

    /// Reloads the Files tab (agent-touched files with badges).
    func refreshWorkspaceState() {
        guard let workspace else { return }
        Task { [weak self] in
            await workspace.refreshScan()
            let files = await workspace.files()
            self?.workspaceFiles = files
        }
    }

    /// Reloads the Audit tab from the workspace's append-only JSONL log.
    func refreshAudit() {
        guard let audit else { return }
        Task { [weak self] in
            let entries = await audit.entries()
            self?.auditEntries = entries
        }
    }

    /// The workspace root, when one exists (Files tab, workspace chip).
    var workspaceRoot: URL? { workspace?.root ?? task?.workspaceURL }

    /// Renames a plan item from the Plan panel. Display-side only: the agent
    /// re-reads the plan on its next `update_plan` call, and the edit is
    /// carried in the plan attached to the run's history entry.
    func renamePlanItem(id: UUID, to title: String) {
        guard var plan else { return }
        let trimmed = title.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !trimmed.isEmpty,
              let index = plan.items.firstIndex(where: { $0.id == id }) else { return }
        plan.items[index].title = trimmed
        plan.updatedAt = Date()
        self.plan = plan
    }

    // MARK: - Helpers

    private func mutateStep(id: UUID, _ mutate: (inout LiveStep) -> Void) {
        for index in entries.indices.reversed() {
            if case .step(var step) = entries[index], step.id == id {
                mutate(&step)
                entries[index] = .step(step)
                return
            }
        }
    }

    private func mutateLastStep(_ mutate: (inout LiveStep) -> Void) {
        for index in entries.indices.reversed() {
            if case .step(var step) = entries[index] {
                mutate(&step)
                entries[index] = .step(step)
                return
            }
        }
    }

    private func updateTaskStatus(_ status: AgentTaskStatus) {
        store.setStatus(status, for: taskID)
    }

    private func appendTerminal(_ kind: TerminalLine.Kind, _ text: String) {
        terminalLines.append(TerminalLine(kind: kind, text: text))
        if terminalLines.count > Self.terminalLineCap {
            terminalLines.removeFirst(terminalLines.count - Self.terminalLineCap)
        }
    }

    private func payloadPreview(of call: ToolCall) -> String? {
        guard let arguments = call.argumentsDictionary else { return call.name }
        for key in ["command", "script", "path"] {
            if let value = arguments[key] as? String {
                return call.name == "bash" ? value : "\(call.name): \(value)"
            }
        }
        return call.name
    }
}

// MARK: - RunHub

/// Keeps one RunController per task alive for the app session, so a run keeps
/// streaming (and stays cancellable) while the user browses other tasks.
@MainActor
final class RunHub: ObservableObject {
    private var controllers: [AgentTask.ID: RunController] = [:]
    private let store: TaskStore
    private let settings: AgentSettingsStore?

    init(store: TaskStore, settings: AgentSettingsStore? = nil) {
        self.store = store
        self.settings = settings
    }

    func controller(for taskID: AgentTask.ID) -> RunController {
        if let existing = controllers[taskID] { return existing }
        let controller = RunController(taskID: taskID, store: store, settings: settings)
        controllers[taskID] = controller
        return controller
    }

    /// Drops the controller for a deleted task (cancelling any live run).
    func remove(taskID: AgentTask.ID) {
        controllers[taskID]?.cancel()
        controllers[taskID] = nil
    }

    /// The controller of any task currently running (menu-bar/⌘. targets).
    var runningControllers: [RunController] {
        controllers.values.filter(\.isRunning)
    }
}