//
//  MemoryManager.swift
//  Zyquo Agent
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//
//  Context compression for long autonomous runs (docs/AGENT-RESEARCH.md §3):
//
//  • Token accounting — chars/4 heuristic, calibrated by the provider's
//    reported input tokens when available; the model's contextWindow bounds it.
//  • Compaction at 85% — older completed steps' messages are replaced by ONE
//    synthetic summary message produced by a summarization call to the SAME
//    provider/model (with a mechanical fallback when that call fails). Always
//    kept verbatim: the system prompt (never in history), the original task
//    message, the most recent N steps, and — re-injected fresh inside the
//    summary — the current plan and MEMORY.md. A thrash guard enforces a
//    minimum step distance between compactions.
//  • Output offloading — tool results above a size threshold are written to
//    `.zyquo/outputs/step<N>-<tool>-<id>.txt` and replaced in-context by a
//    head-of-output stub pointing at the file (readable via read_file /
//    search_files).
//  • MEMORY.md — created in the workspace root at task start; the system
//    prompt gives the model ownership of it via the file tools; its content
//    is re-read and re-injected at each compaction.
//

import Foundation

/// Compaction/offloading tunables (Settings › Agent in Phase 6).
struct MemoryConfiguration: Sendable {
    /// Fraction of the model's context window that triggers compaction.
    var compactionThreshold = 0.85
    /// Steps whose messages are always kept verbatim (most recent).
    var keepRecentSteps = 6
    /// Thrash guard: minimum steps between two compactions.
    var minStepsBetweenCompactions = 10
    /// Tool-result content above this many bytes is offloaded to a file.
    var offloadThresholdBytes = 8_192
    /// Lines of the original output kept in the in-context stub.
    var offloadStubLines = 40
    /// Max tokens requested from the summarization call.
    var summaryMaxTokens = 1_200

    static let `default` = MemoryConfiguration()
}

/// One message of the loop's conversation history, tagged with the step that
/// produced it so compaction can operate on whole steps (preserving
/// tool_use/tool_result pairing, which is per-step by construction).
struct HistoryEntry: Sendable {
    /// 0 = the initial task message; N = messages of step N.
    var stepIndex: Int
    var message: Message
    /// Never compacted (the original task message).
    var pinned: Bool = false
    /// A synthetic summary produced by an earlier compaction (eligible for
    /// re-summarization by the next one).
    var isSummary: Bool = false
}

/// One compaction, for the transcript and the UI.
struct CompactionRecord: Codable, Sendable {
    var id: UUID = UUID()
    var timestamp: Date = Date()
    /// Estimated context tokens immediately before / after.
    var beforeTokens: Int
    var afterTokens: Int
    /// How many step's message groups were folded into the summary, through
    /// which step index.
    var summarizedSteps: Int
    var throughStep: Int
    /// The synthetic summary text injected into context.
    var summary: String
}

/// Owned by the AgentLoop actor (class: mutable state survives async
/// summarization calls without inout-across-await restrictions).
final class MemoryManager {
    let configuration: MemoryConfiguration
    private let workspaceRoot: URL
    private let outputsDirectory: URL
    private(set) var lastCompactionStep = 0
    /// Latest input-token figure reported by the provider (calibrates the
    /// heuristic estimate).
    var lastReportedInputTokens: Int?

    init(workspaceRoot: URL, outputsDirectory: URL, configuration: MemoryConfiguration = .default) {
        self.workspaceRoot = workspaceRoot
        self.outputsDirectory = outputsDirectory
        self.configuration = configuration
    }

    // MARK: - Token accounting

    /// chars/4 heuristic — deliberately conservative and provider-neutral.
    static func estimateTokens(_ text: String) -> Int {
        text.isEmpty ? 0 : max(1, text.count / 4)
    }

    /// Estimated tokens of one message (text + reasoning + tool calls +
    /// tool results + fixed per-message overhead).
    static func estimateTokens(of message: Message) -> Int {
        var total = 8 // role/framing overhead
        total += estimateTokens(message.text)
        if let reasoning = message.reasoning { total += estimateTokens(reasoning) }
        for call in message.toolCalls ?? [] {
            total += estimateTokens(call.name) + estimateTokens(call.argumentsJSON) + 8
        }
        for result in message.toolResults ?? [] {
            total += estimateTokens(result.content) + 8
        }
        return total
    }

    func estimateTokens(entries: [HistoryEntry]) -> Int {
        entries.reduce(0) { $0 + Self.estimateTokens(of: $1.message) }
    }

    /// Fixed per-request overhead: system prompt + tool schemas.
    static func estimateFixedTokens(systemPrompt: String, toolSpecs: [ToolSpec]) -> Int {
        var total = estimateTokens(systemPrompt)
        for spec in toolSpecs {
            total += estimateTokens(spec.name) + estimateTokens(spec.description) + estimateTokens(spec.parametersJSONSchema)
        }
        return total
    }

    /// Best current estimate of the next request's input size.
    func currentContextTokens(entries: [HistoryEntry], fixedTokens: Int) -> Int {
        max(fixedTokens + estimateTokens(entries: entries), lastReportedInputTokens ?? 0)
    }

    /// True when usage crossed the threshold, the thrash guard allows it, and
    /// there is at least one whole old step to fold away.
    func shouldCompact(entries: [HistoryEntry], fixedTokens: Int, contextWindow: Int, currentStep: Int) -> Bool {
        guard contextWindow > 0 else { return false }
        let usage = currentContextTokens(entries: entries, fixedTokens: fixedTokens)
        guard Double(usage) >= configuration.compactionThreshold * Double(contextWindow) else { return false }
        guard currentStep - lastCompactionStep >= configuration.minStepsBetweenCompactions else { return false }
        let cutoff = currentStep - configuration.keepRecentSteps
        return entries.contains { !$0.pinned && $0.stepIndex <= cutoff && $0.stepIndex >= 1 }
    }

    // MARK: - Compaction

    struct CompactionOutcome {
        var entries: [HistoryEntry]
        var record: CompactionRecord
    }

    /// Folds every non-pinned message of steps ≤ (currentStep − keepRecent)
    /// — including previous summaries — into one synthetic user message:
    /// summary + fresh plan + fresh MEMORY.md. Returns nil when there is
    /// nothing to compact.
    func compact(
        entries: [HistoryEntry],
        planContext: String?,
        steps: [AgentStep],
        client: any ProviderClient,
        apiKey: String,
        model: AIModel,
        currentStep: Int,
        fixedTokens: Int
    ) async -> CompactionOutcome? {
        let cutoff = currentStep - configuration.keepRecentSteps
        let toSummarize = entries.filter { !$0.pinned && $0.stepIndex <= cutoff }
        guard !toSummarize.isEmpty else { return nil }
        let kept = entries.filter { $0.pinned || $0.stepIndex > cutoff }

        let beforeTokens = fixedTokens + estimateTokens(entries: entries)

        // Summarize via the same provider/model; mechanical fallback on error.
        let material = Self.renderMaterial(toSummarize)
        var summary: String
        do {
            summary = try await Self.requestSummary(
                material: material,
                client: client,
                apiKey: apiKey,
                model: model,
                maxTokens: configuration.summaryMaxTokens
            )
        } catch {
            summary = Self.mechanicalSummary(steps: steps, throughStep: cutoff)
            summary += "\n(Note: the summarization call failed — this is a mechanical digest. Full history: .zyquo/transcript.json.)"
        }

        var checkpoint = "[CONTEXT CHECKPOINT — steps 1–\(max(cutoff, 1)) were compacted into this summary]\n"
        checkpoint += summary.trimmingCharacters(in: .whitespacesAndNewlines)
        checkpoint += "\n\nCURRENT PLAN:\n\(planContext ?? "(no plan recorded yet)")"
        let memoryText = readMemoryFile()?.trimmingCharacters(in: .whitespacesAndNewlines)
        checkpoint += "\n\nMEMORY.md:\n\((memoryText?.isEmpty == false ? memoryText! : "(empty)"))"

        let summaryEntry = HistoryEntry(
            stepIndex: max(cutoff, 1),
            message: Message(role: .user, text: checkpoint),
            isSummary: true
        )

        // Pinned early entries (the task message) stay first, then the
        // summary, then the recent steps in their original order.
        let earlyPinned = kept.filter { $0.stepIndex <= cutoff }
        let recent = kept.filter { $0.stepIndex > cutoff }
        let newEntries = earlyPinned + [summaryEntry] + recent

        lastCompactionStep = currentStep
        let afterTokens = fixedTokens + estimateTokens(entries: newEntries)
        let record = CompactionRecord(
            beforeTokens: beforeTokens,
            afterTokens: afterTokens,
            summarizedSteps: Set(toSummarize.map(\.stepIndex)).count,
            throughStep: max(cutoff, 1),
            summary: summary
        )
        return CompactionOutcome(entries: newEntries, record: record)
    }

    /// Renders the messages being folded away into summarization material,
    /// capped so the summarization call itself cannot blow the window.
    static func renderMaterial(_ entries: [HistoryEntry]) -> String {
        var lines: [String] = []
        for entry in entries {
            let message = entry.message
            var line = "[step \(entry.stepIndex)] \(message.role.rawValue):"
            if !message.text.isEmpty { line += " \(message.text)" }
            for call in message.toolCalls ?? [] {
                line += "\n  → \(call.name) \(String(call.argumentsJSON.prefix(400)))"
            }
            for result in message.toolResults ?? [] {
                let status = result.isError ? "ERROR" : "ok"
                line += "\n  ← [\(status)] \(String(result.content.prefix(600)))"
            }
            lines.append(line)
        }
        var material = lines.joined(separator: "\n")
        let cap = 60_000
        if material.count > cap {
            // Keep the head (task framing) and the tail (recent detail).
            let head = String(material.prefix(cap / 3))
            let tail = String(material.suffix(cap * 2 / 3))
            material = head + "\n… [middle elided] …\n" + tail
        }
        return material
    }

    /// The summarization call: same provider/model, non-streaming, no tools.
    private static func requestSummary(
        material: String,
        client: any ProviderClient,
        apiKey: String,
        model: AIModel,
        maxTokens: Int
    ) async throws -> String {
        let instruction = """
            You compress an autonomous agent's working history into a compact \
            structured record. From the transcript below, produce EXACTLY \
            these three sections, tersely, preserving concrete values \
            (paths, filenames, commands, numbers, error messages):

            Sub-tasks completed:
            Key facts, paths & decisions:
            Open issues:

            TRANSCRIPT:
            \(material)
            """
        var parameters = ChatParameters(maxTokens: maxTokens)
        parameters.temperature = nil
        let request = ChatRequest(
            model: model,
            systemPrompt: "You are a precise summarizer inside the Zyquo Agent memory system. Output only the requested record — no preamble.",
            messages: [Message(role: .user, text: instruction)],
            parameters: parameters,
            stream: false
        )
        let reply = try await client.complete(request, apiKey: apiKey)
        let text = reply.text.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !text.isEmpty else {
            throw ProviderError.invalidResponse(client.providerID, detail: "empty summarization reply")
        }
        return text
    }

    /// Fallback digest built from the transcript's AgentSteps when the
    /// summarization call fails (offline, rate-limited, mock runs).
    static func mechanicalSummary(steps: [AgentStep], throughStep: Int) -> String {
        var lines = ["Sub-tasks completed:"]
        for step in steps where step.index >= 1 && step.index <= throughStep {
            var line = "- Step \(step.index):"
            let thought = step.text.trimmingCharacters(in: .whitespacesAndNewlines)
            if !thought.isEmpty { line += " \(String(thought.prefix(120)))" }
            if !step.toolInvocations.isEmpty {
                let calls = step.toolInvocations.map { invocation -> String in
                    let ok = invocation.result?.isError == true ? "error" : "ok"
                    return "\(invocation.call.name)(\(String(invocation.call.argumentsJSON.prefix(60))))→\(ok)"
                }
                line += " [\(calls.joined(separator: ", "))]"
            }
            lines.append(line)
        }
        lines.append("Key facts, paths & decisions:\n- (not extracted — mechanical digest)")
        lines.append("Open issues:\n- (unknown)")
        return lines.joined(separator: "\n")
    }

    // MARK: - Output offloading

    /// Writes oversized tool-result content to `.zyquo/outputs/…` and returns
    /// a stub result (head of the output + pointer). Small results pass
    /// through untouched, as do failures to write the file.
    func offloadIfNeeded(_ result: ToolResult, stepIndex: Int, toolName: String) -> ToolResult {
        guard result.content.utf8.count > configuration.offloadThresholdBytes else { return result }

        let safeTool = toolName.map { $0.isLetter || $0.isNumber ? $0 : "-" }
        let safeID = result.toolCallID.suffix(8).map { $0.isLetter || $0.isNumber ? $0 : "-" }
        let fileName = "step\(stepIndex)-\(String(safeTool))-\(String(safeID)).txt"
        let fileURL = outputsDirectory.appendingPathComponent(fileName)

        do {
            try FileManager.default.createDirectory(at: outputsDirectory, withIntermediateDirectories: true)
            try result.content.write(to: fileURL, atomically: true, encoding: .utf8)
        } catch {
            // Offloading is an optimization — never fail the tool call over it.
            return result
        }

        let allLines = result.content.components(separatedBy: "\n")
        var head = allLines.prefix(configuration.offloadStubLines).joined(separator: "\n")
        if head.count > 4_000 { head = String(head.prefix(4_000)) }
        let sizeKB = Double(result.content.utf8.count) / 1000
        let relativePath = ".zyquo/outputs/\(fileName)"
        let stub = head + """
            \n… [output truncated: \(String(format: "%.1f", sizeKB)) KB, \(allLines.count) lines total — \
            full output: \(relativePath) — use read_file or search_files on that path]
            """
        return ToolResult(toolCallID: result.toolCallID, content: stub, isError: result.isError)
    }

    // MARK: - MEMORY.md

    var memoryFileURL: URL { workspaceRoot.appendingPathComponent("MEMORY.md") }

    /// Creates MEMORY.md with its contract header if absent (task start).
    func ensureMemoryFile() {
        guard !FileManager.default.fileExists(atPath: memoryFileURL.path) else { return }
        let template = """
            # MEMORY.md — Zyquo Agent task memory

            Durable facts, decisions, and open questions for this task.
            Maintained by the agent (see system prompt): read at task start,
            append important knowledge as it is learned. Content here survives
            context compaction.

            """
        try? template.write(to: memoryFileURL, atomically: true, encoding: .utf8)
    }

    /// Current MEMORY.md content (capped — memory is notes, not a dump).
    func readMemoryFile() -> String? {
        guard let text = try? String(contentsOf: memoryFileURL, encoding: .utf8) else { return nil }
        return text.count > 16_000 ? String(text.prefix(16_000)) + "\n… [MEMORY.md truncated]" : text
    }
}