//
//  Distiller.swift
//  Prisme
//
//  Author: Simon-Pierre Boucher <contact@spboucher.ai>
//

import Foundation
import FoundationModels

/// HTML → compact structure (CLAUDE.md §4, the heart of the project).
/// Steps 1–2 (extraction, typed blocks) happen in the page via extractor.js.
/// This actor performs step 3 (token budgeting), a deterministic form of
/// step 4 (condensation), and step 5 (guided generation). Hierarchical
/// LLM condensation of low-priority sections comes later — the budgeter
/// currently trims deterministically, never mid-sentence.
actor Distiller {
    enum DistillerError: LocalizedError {
        case emptyPage
        case overBudget

        var errorDescription: String? {
            switch self {
            case .emptyPage: "La page ne contient pas assez de texte lisible."
            case .overBudget: "La page dépasse la capacité du modèle local."
            }
        }
    }

    private let cache = DigestCache()

    /// Fraction of the context window reserved for the model's response —
    /// a prompt that fills the window leaves no room to answer (§2).
    private static let responseReserve = 0.30

    func digest(
        _ content: ExtractedContent,
        queue: InferenceQueue,
        priority: InferencePriority
    ) async throws -> (PageDigest, [ContentBlock]) {
        guard !content.blocks.isEmpty else { throw DistillerError.emptyPage }

        let key = DigestCache.key(for: content)
        if let hit = await cache.lookup(key) {
            return hit
        }

        let model = SystemLanguageModel.default
        let instructions = """
        You are the reading engine of a browser. You receive a numbered list \
        of text blocks extracted from one web page. Produce the requested \
        structured digest using ONLY the provided text — never outside \
        knowledge, never invented facts. If the page does not support a \
        field, keep it minimal. Write summaries in the language of the page. \
        Every sourceBlock must be the index of a block from the list.
        """

        let (prompt, keptBlocks) = try await budgetedPrompt(
            content: content,
            model: model,
            instructionsCost: try await model.tokenCount(for: instructions)
        )

        let digest = try await queue.run(priority) {
            let session = LanguageModelSession(model: model, instructions: instructions)
            let response = try await session.respond(to: prompt, generating: PageDigest.self)
            return response.content
        }

        await cache.store(digest, blocks: keptBlocks, key: key)
        return (digest, keptBlocks)
    }

    /// Step 3: measure and allocate. Headings are always kept (they are the
    /// skeleton); body blocks are added by document order until the budget
    /// is spent, long blocks trimmed at a sentence boundary.
    private func budgetedPrompt(
        content: ExtractedContent,
        model: SystemLanguageModel,
        instructionsCost: Int
    ) async throws -> (String, [ContentBlock]) {
        let contextSize = try await model.contextSize
        let budget = Int(Double(contextSize) * (1 - Self.responseReserve)) - instructionsCost - 200
        guard budget > 300 else { throw DistillerError.overBudget }

        var lines: [String] = ["Page title: \(content.title)", "Blocks:"]
        var kept: [ContentBlock] = []
        var spent = try await model.tokenCount(for: lines.joined(separator: "\n"))

        for block in content.blocks {
            var text = block.text
            if text.count > 600 {
                text = trimToSentence(text, limit: 600)
            }
            let line = "[\(kept.count)] (\(block.kind.rawValue)) \(text)"
            let cost = try await model.tokenCount(for: line)
            if spent + cost > budget {
                if block.kind == .heading {
                    // Headings squeeze in with a shorter form when possible.
                    let short = "[\(kept.count)] (heading) \(String(text.prefix(80)))"
                    let shortCost = try await model.tokenCount(for: short)
                    guard spent + shortCost <= budget else { break }
                    lines.append(short)
                    kept.append(block)
                    spent += shortCost
                    continue
                }
                continue
            }
            lines.append(line)
            kept.append(block)
            spent += cost
        }

        guard !kept.isEmpty else { throw DistillerError.overBudget }
        return (lines.joined(separator: "\n"), kept)
    }

    private func trimToSentence(_ text: String, limit: Int) -> String {
        guard text.count > limit else { return text }
        let head = String(text.prefix(limit))
        if let cut = head.lastIndex(where: { ".!?".contains($0) }) {
            return String(head[...cut])
        }
        if let space = head.lastIndex(of: " ") {
            return String(head[..<space]) + "…"
        }
        return head + "…"
    }
}