//
//  LoopGuard.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Bounds every run (docs/AGENT-RESEARCH.md §6.1): max steps, cumulative
//  token budget, wall-clock budget, repetition detection (the same tool
//  call failing N times), and stall detection (M consecutive steps with no
//  plan change, no file change, and no new distinct tool call). A trip never
//  aborts the run — the loop pauses and hands control back to the user
//  ("pause, ask" is the correct trip behavior, per Anthropic's own
//  pause_turn design), who can continue with raised budgets or stop.
//

import Foundation

/// Budgets and detection thresholds, user-tunable in Settings › Agent.
struct LoopGuardConfiguration: Sendable {
    /// Hard cap on loop iterations per run.
    var maxSteps = 50
    /// Cumulative token budget (input + output + reasoning) per run.
    var tokenBudget = 500_000
    /// Wall-clock budget per run.
    var wallClockBudget: TimeInterval = 30 * 60
    /// Same tool + normalized args failing this many times ⇒ repetition trip.
    var repetitionThreshold = 3
    /// Consecutive no-progress steps ⇒ stall trip.
    var stallThreshold = 6

    static let `default` = LoopGuardConfiguration()
}

/// Why the guard tripped.
enum LoopGuardTripReason: String, Codable, Sendable {
    case maxSteps
    case tokenBudget
    case wallClockBudget
    case repetition
    case stall
}

/// One trip, surfaced to the UI/CLI so the user can decide.
struct LoopGuardTrip: Codable, Sendable {
    var reason: LoopGuardTripReason
    /// Human-readable explanation shown on the pause card.
    var message: String
    /// Step at which the trip occurred.
    var stepIndex: Int
}

/// Run-bounding state machine. Plain struct — mutated only inside the
/// AgentLoop actor.
struct LoopGuard {
    private(set) var configuration: LoopGuardConfiguration
    private let startedAt = Date()
    private(set) var totalTokens = 0

    /// Consecutive failure count per tool-call signature (reset on success).
    private var failingCallCounts: [String: Int] = [:]
    /// Every distinct tool-call signature seen this run (novelty = progress).
    private var seenCallSignatures: Set<String> = []
    private var stepsWithoutProgress = 0

    init(configuration: LoopGuardConfiguration = .default) {
        self.configuration = configuration
    }

    var elapsed: TimeInterval { Date().timeIntervalSince(startedAt) }

    // MARK: Budgets (checked before each step)

    /// Nil when the next step may proceed; a trip otherwise.
    mutating func checkBeforeStep(index: Int) -> LoopGuardTrip? {
        if index > configuration.maxSteps {
            return LoopGuardTrip(
                reason: .maxSteps,
                message: "Reached the step limit (\(configuration.maxSteps) steps).",
                stepIndex: index
            )
        }
        if totalTokens >= configuration.tokenBudget {
            return LoopGuardTrip(
                reason: .tokenBudget,
                message: "Reached the token budget (\(totalTokens) of \(configuration.tokenBudget) tokens used).",
                stepIndex: index
            )
        }
        if elapsed >= configuration.wallClockBudget {
            return LoopGuardTrip(
                reason: .wallClockBudget,
                message: "Reached the time budget (\(Int(elapsed / 60)) of \(Int(configuration.wallClockBudget / 60)) minutes).",
                stepIndex: index
            )
        }
        return nil
    }

    mutating func recordUsage(_ usage: TokenUsage) {
        totalTokens += usage.totalTokens + (usage.reasoningTokens ?? 0)
    }

    // MARK: Repetition & novelty (recorded per tool invocation)

    /// Records one executed tool call. Returns true when the call was a NEW
    /// distinct signature (a progress signal for stall detection).
    mutating func recordInvocation(call: ToolCall, isError: Bool) -> Bool {
        let signature = Self.signature(of: call)
        let isNew = seenCallSignatures.insert(signature).inserted
        if isError {
            failingCallCounts[signature, default: 0] += 1
        } else {
            failingCallCounts[signature] = 0
        }
        return isNew
    }

    // MARK: Post-step verdict

    /// Checks repetition and stall after a tool-executing step. `planChanged`,
    /// `filesChanged`, and `sawNewToolCall` are the progress signals.
    mutating func recordStepOutcome(
        stepIndex: Int,
        planChanged: Bool,
        filesChanged: Bool,
        sawNewToolCall: Bool
    ) -> LoopGuardTrip? {
        if let (signature, count) = failingCallCounts.first(where: { $0.value >= configuration.repetitionThreshold }) {
            // Reset so a user-resumed run isn't instantly re-tripped.
            failingCallCounts[signature] = 0
            let display = signature.count > 160 ? String(signature.prefix(160)) + "…" : signature
            return LoopGuardTrip(
                reason: .repetition,
                message: "The same tool call has failed \(count) times: \(display)",
                stepIndex: stepIndex
            )
        }

        if planChanged || filesChanged || sawNewToolCall {
            stepsWithoutProgress = 0
        } else {
            stepsWithoutProgress += 1
        }
        if stepsWithoutProgress >= configuration.stallThreshold {
            stepsWithoutProgress = 0
            return LoopGuardTrip(
                reason: .stall,
                message: "No progress detected for \(configuration.stallThreshold) consecutive steps (no plan change, no file change, no new tool call).",
                stepIndex: stepIndex
            )
        }
        return nil
    }

    // MARK: Resume support

    /// Raises every budget by half its current value (with sensible floors)
    /// and clears the repetition/stall state — called when the user answers
    /// a trip with "continue".
    mutating func raiseBudgets() {
        configuration.maxSteps += max(10, configuration.maxSteps / 2)
        configuration.tokenBudget += max(100_000, configuration.tokenBudget / 2)
        configuration.wallClockBudget += max(600, configuration.wallClockBudget / 2)
        failingCallCounts = [:]
        stepsWithoutProgress = 0
    }

    // MARK: Signature normalization

    /// "name|canonical-args": arguments parsed and re-serialized with sorted
    /// keys so cosmetic JSON differences don't defeat repetition detection.
    static func signature(of call: ToolCall) -> String {
        "\(call.name)|\(normalizedArguments(call.argumentsJSON))"
    }

    static func normalizedArguments(_ argumentsJSON: String) -> String {
        guard let value = JSONValue.parse(argumentsJSON) else {
            return argumentsJSON.trimmingCharacters(in: .whitespacesAndNewlines)
        }
        let encoder = JSONEncoder()
        encoder.outputFormatting = [.sortedKeys]
        guard let data = try? encoder.encode(value),
              let canonical = String(data: data, encoding: .utf8) else {
            return argumentsJSON.trimmingCharacters(in: .whitespacesAndNewlines)
        }
        return canonical
    }
}