//
//  AgentLoop.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  The plan→act→observe→reflect engine (docs/AGENT-RESEARCH.md §1.3): a
//  while-loop keyed on the normalized StopReason. Each iteration streams one
//  model turn (text, thinking, and tool-call arguments live), executes any
//  tool calls through the ToolRegistry — every side effect pre-cleared by
//  the PolicyEngine and audited — threads the results back as one message
//  correlated by call id, offloads oversized outputs, compacts memory near
//  the context limit, and checks the LoopGuard. A turn without tool calls
//  is the done-signal: its text is the user-facing result. Trips pause the
//  run (resume/stop), they never abort it; cancellation kills any in-flight
//  process via Task cancellation.
//

import Foundation

/// Per-run engine tunables (Settings › Agent in Phase 6).
struct AgentConfiguration: Sendable {
    var loopGuard = LoopGuardConfiguration.default
    var memory = MemoryConfiguration.default
    /// Execute a turn's tool calls concurrently when they are ALL read-only
    /// under the active policy. Mutating calls always run sequentially.
    var parallelToolCalls = false
    /// Default per-command timeout for the ExecutionService the host wires up
    /// (the loop itself never spawns processes).
    var perCommandTimeout: TimeInterval?
    /// Active persona's system-prompt addition, appended to the agent system
    /// prompt as a trailing "## Persona" section (Phase 6 personas).
    var personaAddendum: String?

    static let `default` = AgentConfiguration()
}

/// The agent engine. Actor: one run at a time, all mutable run state
/// (history, plan, guard, transcript) confined here.
actor AgentLoop {
    // MARK: Dependencies

    private let model: AIModel
    private let client: any ProviderClient
    private let apiKey: String
    private let tools: ToolRegistry
    private let policy: PolicyEngine
    private let audit: AuditLog
    private let workspace: WorkspaceManager
    private let configuration: AgentConfiguration

    // MARK: Run state

    private var planner: Planner
    private var loopGuard: LoopGuard
    private let memory: MemoryManager
    private var transcript: Transcript?
    private var history: [HistoryEntry] = []
    private var runTask: Task<Void, Never>?
    private var eventContinuation: AsyncThrowingStream<AgentEvent, Error>.Continuation?
    private var tripContinuation: CheckedContinuation<TripDecision, Never>?
    private var isRunning = false

    /// The user's answer to a LoopGuard trip.
    enum TripDecision: Sendable {
        case resume(raiseBudget: Bool)
        case stop
    }

    init(
        model: AIModel,
        client: any ProviderClient,
        apiKey: String,
        tools: ToolRegistry,
        policy: PolicyEngine,
        audit: AuditLog,
        workspace: WorkspaceManager,
        configuration: AgentConfiguration = .default
    ) {
        self.model = model
        self.client = client
        self.apiKey = apiKey
        self.tools = tools
        self.policy = policy
        self.audit = audit
        self.workspace = workspace
        self.configuration = configuration
        self.planner = Planner(workspaceRoot: workspace.root)
        self.loopGuard = LoopGuard(configuration: configuration.loopGuard)
        self.memory = MemoryManager(
            workspaceRoot: workspace.root,
            outputsDirectory: workspace.outputsDirectory,
            configuration: configuration.memory
        )
    }

    // MARK: - Public API

    /// Starts the run and returns its live event stream. The stream ends
    /// after `.runFinished`; provider failures that end the run are reported
    /// through `.runFinished(.failed…)`, not thrown.
    func run(task: String) -> AsyncThrowingStream<AgentEvent, Error> {
        let (stream, continuation) = AsyncThrowingStream<AgentEvent, Error>.makeStream()
        guard !isRunning else {
            continuation.yield(.runFinished(.failed(reason: "A run is already in progress on this AgentLoop.")))
            continuation.finish()
            return stream
        }
        isRunning = true
        eventContinuation = continuation
        runTask = Task {
            await self.execute(task: task)
        }
        return stream
    }

    /// Stops the run: cancels the in-flight model stream and any executing
    /// process (ExecutionService honors Task cancellation with SIGTERM→SIGKILL).
    func cancel() {
        runTask?.cancel()
        // A run paused on a guard trip is not awaiting a cancellable
        // suspension — unblock it explicitly.
        tripContinuation?.resume(returning: .stop)
        tripContinuation = nil
    }

    /// Answers a `.guardTripped` pause: continue, optionally raising every
    /// budget (without raising, a budget trip will re-trip on the next step;
    /// repetition/stall trips reset their counters either way).
    func resume(raisingBudget: Bool = true) {
        guard let continuation = tripContinuation else { return }
        tripContinuation = nil
        continuation.resume(returning: .resume(raiseBudget: raisingBudget))
    }

    /// Answers a `.guardTripped` pause: end the run.
    func stop() {
        guard let continuation = tripContinuation else { return }
        tripContinuation = nil
        continuation.resume(returning: .stop)
    }

    // MARK: - The loop

    private func execute(task: String) async {
        emit(.statusChanged(.preparing))

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

        memory.ensureMemoryFile()
        transcript = Transcript(task: task, model: model, workspaceRoot: workspace.root)
        history = [HistoryEntry(stepIndex: 0, message: Message(role: .user, text: task), pinned: true)]
        if let plan = planner.plan {
            // Reopened workspace with a persisted plan — surface it.
            emit(.planUpdated(plan))
        }

        let systemPrompt = AgentSystemPrompt.build(
            workspacePath: workspace.root.path,
            toolNames: tools.toolNames,
            safetyMode: await policy.mode,
            personaAddendum: configuration.personaAddendum
        )
        let toolSpecs = tools.toolSpecs + [Planner.toolSpec]
        let fixedTokens = MemoryManager.estimateFixedTokens(systemPrompt: systemPrompt, toolSpecs: toolSpecs)

        var stepIndex = 0
        var truncationNudges = 0
        var emptyAnswerNudges = 0

        while true {
            if Task.isCancelled {
                finish(.cancelled)
                return
            }
            stepIndex += 1

            // ---- LoopGuard budgets (pause-and-ask, never abort) ----------
            if let trip = loopGuard.checkBeforeStep(index: stepIndex) {
                if await pauseOnTrip(trip) == false {
                    finish(.stoppedByUser(reason: trip.message))
                    return
                }
            }

            // ---- Memory compaction (before the request, at step boundary) -
            if memory.shouldCompact(entries: history, fixedTokens: fixedTokens, contextWindow: model.contextWindow, currentStep: stepIndex - 1) {
                emit(.statusChanged(.compacting))
                if let outcome = await memory.compact(
                    entries: history,
                    planContext: planner.renderedForContext(),
                    steps: transcript?.document.steps ?? [],
                    client: client,
                    apiKey: apiKey,
                    model: model,
                    currentStep: stepIndex - 1,
                    fixedTokens: fixedTokens
                ) {
                    history = outcome.entries
                    transcript?.recordCompaction(outcome.record)
                    emit(.compactionPerformed(outcome.record))
                }
            }

            // ---- One model turn, streamed --------------------------------
            var step = AgentStep(index: stepIndex, text: "")
            emit(.stepStarted(step))
            emit(.statusChanged(.running))

            let turn: ModelTurn
            do {
                turn = try await streamModelTurn(stepID: step.id, systemPrompt: systemPrompt, toolSpecs: toolSpecs)
            } catch is CancellationError {
                finish(.cancelled)
                return
            } catch {
                step.status = .failed
                step.finishedAt = Date()
                transcript?.upsert(step: step)
                emit(.stepCompleted(step))
                finish(.failed(reason: error.localizedDescription))
                return
            }

            step.text = turn.text
            step.thinking = turn.reasoning.isEmpty ? nil : turn.reasoning
            if let usage = turn.usage {
                step.inputTokens = usage.inputTokens
                step.outputTokens = usage.outputTokens
                loopGuard.recordUsage(usage)
                transcript?.addUsage(usage)
                memory.lastReportedInputTokens = usage.inputTokens
            }

            // ---- Branch on the normalized stop reason --------------------
            switch turn.stop {
            case .toolUse where !turn.toolCalls.isEmpty:
                history.append(HistoryEntry(
                    stepIndex: stepIndex,
                    message: .assistantToolCalls(turn.toolCalls, text: turn.text, reasoning: step.thinking)
                ))
                step.status = .executing
                emit(.statusChanged(.executingTools))

                let filesBefore = await filesSignature()
                let execution: ToolExecutionOutcome
                do {
                    execution = try await executeToolCalls(turn.toolCalls, stepIndex: stepIndex, stepID: step.id)
                } catch {
                    // Only CancellationError escapes the registry.
                    step.status = .cancelled
                    step.finishedAt = Date()
                    transcript?.upsert(step: step)
                    emit(.stepCompleted(step))
                    finish(.cancelled)
                    return
                }
                history.append(HistoryEntry(stepIndex: stepIndex, message: .toolResultsMessage(execution.results)))
                await workspace.refreshScan()
                let filesAfter = await filesSignature()

                step.toolInvocations = execution.invocations
                step.status = .completed
                step.finishedAt = Date()
                transcript?.upsert(step: step)
                emit(.stepCompleted(step))

                if let trip = loopGuard.recordStepOutcome(
                    stepIndex: stepIndex,
                    planChanged: execution.planChanged,
                    filesChanged: filesBefore != filesAfter,
                    sawNewToolCall: execution.sawNewToolCall
                ) {
                    if await pauseOnTrip(trip) == false {
                        finish(.stoppedByUser(reason: trip.message))
                        return
                    }
                }

            case .endTurn, .toolUse, .other:
                // .toolUse with an empty call list and .other are anomalies —
                // treat a non-empty text as the final answer, nudge otherwise.
                let answer = turn.text.trimmingCharacters(in: .whitespacesAndNewlines)
                if answer.isEmpty, stepIndex > 1, emptyAnswerNudges < 1 {
                    emptyAnswerNudges += 1
                    step.status = .completed
                    step.finishedAt = Date()
                    transcript?.upsert(step: step)
                    emit(.stepCompleted(step))
                    history.append(HistoryEntry(
                        stepIndex: stepIndex,
                        message: Message(role: .user, text: "Your last message was empty. If the task is complete, state the result summary; otherwise continue with the next tool call.")
                    ))
                    continue
                }
                // A stream cancelled mid-turn surfaces here as an empty
                // "end turn"; that is a cancelled run, not a completed one.
                if Task.isCancelled {
                    step.status = .cancelled
                    step.finishedAt = Date()
                    transcript?.upsert(step: step)
                    emit(.stepCompleted(step))
                    finish(.cancelled)
                    return
                }
                step.status = .completed
                step.finishedAt = Date()
                transcript?.upsert(step: step)
                emit(.stepCompleted(step))
                finish(.completed(finalAnswer: turn.text))
                return

            case .maxTokens:
                // A truncated turn may carry INCOMPLETE tool calls — never
                // execute them (docs/AGENT-RESEARCH.md §6.2). Nudge instead.
                truncationNudges += 1
                step.status = .completed
                step.finishedAt = Date()
                transcript?.upsert(step: step)
                emit(.stepCompleted(step))
                if truncationNudges > 2 {
                    finish(.failed(reason: "The model's output was truncated by its token limit three times in a row — the task cannot proceed. Try a model with a larger output limit or split the task."))
                    return
                }
                if !turn.text.isEmpty {
                    history.append(HistoryEntry(stepIndex: stepIndex, message: Message(role: .assistant, text: turn.text)))
                }
                history.append(HistoryEntry(
                    stepIndex: stepIndex,
                    message: Message(role: .user, text: "Your previous response was cut off by the output token limit; any tool calls in it were NOT executed. Continue more concisely, re-issuing the needed tool call(s) one at a time.")
                ))

            case .refusal:
                step.status = .failed
                step.finishedAt = Date()
                transcript?.upsert(step: step)
                emit(.stepCompleted(step))
                finish(.failed(reason: "The model declined to continue this task\(turn.text.isEmpty ? "." : ": \(turn.text)")"))
                return
            }
        }
    }

    // MARK: - Model turn streaming

    /// One accumulated assistant turn.
    private struct ModelTurn {
        var text = ""
        var reasoning = ""
        var toolCalls: [ToolCall] = []
        var usage: TokenUsage?
        var stop = StopReason.endTurn
    }

    /// Streams one turn, forwarding deltas as AgentEvents. Retries once with
    /// backoff on transient provider errors (network, 429, 5xx).
    private func streamModelTurn(stepID: UUID, systemPrompt: String, toolSpecs: [ToolSpec]) async throws -> ModelTurn {
        var attempt = 0
        while true {
            attempt += 1
            do {
                return try await streamOnce(stepID: stepID, systemPrompt: systemPrompt, toolSpecs: toolSpecs)
            } catch let error as ProviderError where attempt == 1 && Self.isTransient(error) {
                let delay: TimeInterval
                if case .rateLimited(_, let retryAfter) = error, let retryAfter {
                    delay = min(retryAfter, 30)
                } else {
                    delay = 2
                }
                try await Task.sleep(nanoseconds: UInt64(delay * 1_000_000_000))
            }
        }
    }

    private static func isTransient(_ error: ProviderError) -> Bool {
        switch error {
        case .networkError, .rateLimited:
            return true
        case .serverError(_, let status, _):
            return status >= 500
        default:
            return false
        }
    }

    private func streamOnce(stepID: UUID, systemPrompt: String, toolSpecs: [ToolSpec]) async throws -> ModelTurn {
        let request = ChatRequest(
            model: model,
            systemPrompt: systemPrompt,
            messages: history.map(\.message),
            parameters: ChatParameters(),
            stream: true,
            tools: toolSpecs,
            toolChoice: .auto
        )
        var turn = ModelTurn()
        for try await event in client.streamChat(request, apiKey: apiKey) {
            switch event {
            case .textDelta(let delta):
                turn.text += delta
                emit(.textDelta(stepID: stepID, delta))
            case .reasoningDelta(let delta):
                turn.reasoning += delta
                emit(.thinkingDelta(stepID: stepID, delta))
            case .toolCallStarted(let index, let id, let name):
                emit(.toolCallStreaming(stepID: stepID, index: index, id: id, name: name))
            case .toolCallArgumentsDelta(let index, let delta):
                emit(.toolCallArgumentsDelta(stepID: stepID, index: index, delta: delta))
            case .toolCalls(let calls):
                turn.toolCalls = calls
            case .usage(let usage):
                turn.usage = usage
            case .citations:
                break
            case .finished(_, let stop):
                turn.stop = stop
            }
        }
        return turn
    }

    // MARK: - Tool execution

    private struct ToolExecutionOutcome {
        var results: [ToolResult] = []
        var invocations: [AgentToolInvocation] = []
        var planChanged = false
        var sawNewToolCall = false
    }

    /// Executes a turn's tool calls — `update_plan` intercepted internally
    /// (no side effect ⇒ no policy gate), everything else through the
    /// registry (which routes side effects through the PolicyEngine). Results
    /// come back in call order, correlated by id, in ONE tool-results
    /// message. Sequential by default; concurrent only when enabled AND every
    /// call is read-only under the active policy. Throws only on cancellation.
    private func executeToolCalls(
        _ calls: [ToolCall],
        stepIndex: Int,
        stepID: UUID
    ) async throws -> ToolExecutionOutcome {
        var outcome = ToolExecutionOutcome()

        if configuration.parallelToolCalls, calls.count > 1, await allReadOnly(calls) {
            // Concurrent read-only execution; results re-ordered by call index.
            var startedAt: [String: Date] = [:]
            for call in calls {
                let invocation = AgentToolInvocation(call: call, startedAt: Date())
                startedAt[call.id] = invocation.startedAt
                emit(.toolCallStarted(stepID: stepID, invocation))
            }
            let registry = tools
            let collected: [(Int, ToolResult)] = try await withThrowingTaskGroup(of: (Int, ToolResult).self) { group in
                for (index, call) in calls.enumerated() {
                    let context = makeContext(stepID: stepID, callID: call.id)
                    group.addTask {
                        (index, try await registry.execute(call: call, context: context))
                    }
                }
                var results: [(Int, ToolResult)] = []
                for try await item in group { results.append(item) }
                return results.sorted { $0.0 < $1.0 }
            }
            for (index, rawResult) in collected {
                let call = calls[index]
                let result = memory.offloadIfNeeded(rawResult, stepIndex: stepIndex, toolName: call.name)
                outcome.sawNewToolCall = loopGuard.recordInvocation(call: call, isError: result.isError) || outcome.sawNewToolCall
                var invocation = AgentToolInvocation(call: call, result: result, startedAt: startedAt[call.id])
                invocation.finishedAt = Date()
                outcome.results.append(result)
                outcome.invocations.append(invocation)
                emit(.toolCallFinished(stepID: stepID, invocation))
            }
            return outcome
        }

        for call in calls {
            try Task.checkCancellation()
            var invocation = AgentToolInvocation(call: call, startedAt: Date())
            emit(.toolCallStarted(stepID: stepID, invocation))

            let result: ToolResult
            if call.name == Planner.toolName {
                let applied = planner.apply(call: call)
                result = applied.result
                outcome.planChanged = outcome.planChanged || applied.changed
                if let plan = planner.plan {
                    emit(.planUpdated(plan))
                    transcript?.recordPlan(plan, stepIndex: stepIndex)
                }
            } else {
                let context = makeContext(stepID: stepID, callID: call.id)
                let raw = try await tools.execute(call: call, context: context)
                result = memory.offloadIfNeeded(raw, stepIndex: stepIndex, toolName: call.name)
            }

            outcome.sawNewToolCall = loopGuard.recordInvocation(call: call, isError: result.isError) || outcome.sawNewToolCall
            invocation.result = result
            invocation.finishedAt = Date()
            outcome.results.append(result)
            outcome.invocations.append(invocation)
            emit(.toolCallFinished(stepID: stepID, invocation))
        }
        return outcome
    }

    private func makeContext(stepID: UUID, callID: String) -> ToolExecutionContext {
        let continuation = eventContinuation
        return ToolExecutionContext(
            workspaceURL: workspace.root,
            policy: policy,
            audit: audit,
            onOutput: { chunk in
                continuation?.yield(.toolOutput(stepID: stepID, callID: callID, chunk))
            },
            workspaceManager: workspace
        )
    }

    /// True when every call in the batch is safe to run concurrently: the
    /// read-only file tools, or a bash command the policy auto-allows under
    /// the current mode (i.e. its whole payload is read-only). Anything
    /// mutating, plan-touching, or approval-bound runs sequentially.
    private func allReadOnly(_ calls: [ToolCall]) async -> Bool {
        let readOnlyTools: Set<String> = ["read_file", "list_dir", "search_files"]
        for call in calls {
            if readOnlyTools.contains(call.name) { continue }
            if call.name == "bash",
               let command = call.argumentsDictionary?["command"] as? String {
                let ruling = await policy.evaluate(ActionRequest(
                    kind: .shellCommand, payload: command, cwd: workspace.root, explanation: nil
                ))
                if case .allow = ruling { continue }
            }
            return false
        }
        return true
    }

    // MARK: - Progress & pausing

    /// Cheap signature of the workspace's tracked files (stall detection).
    private func filesSignature() async -> Int {
        var hasher = Hasher()
        for entry in await workspace.files() {
            hasher.combine(entry.path)
            hasher.combine(entry.lastTouched)
        }
        return hasher.finalize()
    }

    /// Emits the trip, pauses until resume()/stop()/cancel() answers, and
    /// returns true when the run should continue (with raised budgets).
    private func pauseOnTrip(_ trip: LoopGuardTrip) async -> Bool {
        emit(.guardTripped(trip))
        emit(.statusChanged(.awaitingUser))
        let decision = await withCheckedContinuation { (continuation: CheckedContinuation<TripDecision, Never>) in
            tripContinuation = continuation
        }
        switch decision {
        case .resume(let raiseBudget):
            if raiseBudget {
                loopGuard.raiseBudgets()
            }
            emit(.statusChanged(.running))
            return true
        case .stop:
            return false
        }
    }

    // MARK: - Finishing

    private func emit(_ event: AgentEvent) {
        eventContinuation?.yield(event)
    }

    private func finish(_ outcome: AgentRunOutcome) {
        transcript?.finish(outcome: outcome)
        emit(.statusChanged(.finished))
        emit(.runFinished(outcome))
        eventContinuation?.finish()
        eventContinuation = nil
        isRunning = false
    }
}