import FoundationModels
 
/// Layer 1 — the model's own context: a sliding window that collapses
/// into a rolling summary instead of silently dropping older turns.
struct ConversationWindow {
    private let maxTokens: Int
    private(set) var turns: [ChatTurn] = []
    private(set) var rollingSummary: String?
 
    init(maxTokens: Int = 3200) {
        self.maxTokens = maxTokens
    }
 
    /// Rough heuristic: ~4 characters per token (good enough for budgeting,
    /// not for billing — on-device models don't charge per token anyway).
    private func estimatedTokens(_ text: String) -> Int {
        text.count / 4
    }
 
    private var currentTokens: Int {
        turns.reduce(rollingSummary.map(estimatedTokens) ?? 0) {
            $0 + estimatedTokens($1.text)
        }
    }
 
    mutating func append(_ turn: ChatTurn) async {
        turns.append(turn)
        guard currentTokens > maxTokens, turns.count > 4 else { return }
        await collapseOldest()
    }
 
    /// Move the oldest half of the window into a rolling summary and keep
    /// only the freshest turns verbatim. This is the difference between
    /// "forgetting" and "compressing" — the model still knows the gist.
    private mutating func collapseOldest() async {
        let cut = turns.count / 2
        let evicted = Array(turns.prefix(cut))
        turns.removeFirst(cut)
        rollingSummary = await summarize(evicted, previous: rollingSummary)
    }
 
    private func summarize(_ evicted: [ChatTurn], previous: String?) async -> String {
        guard case .available = SystemLanguageModel.default.availability else {
            return previous ?? ""
        }
        let session = LanguageModelSession(instructions: """
        Summarize the earlier part of this conversation in 2-3 neutral \
        sentences. Keep names, decisions and open questions; drop small talk.
        """)
        let transcript = evicted.map { "\($0.role): \($0.text)" }.joined(separator: "\n")
        let prompt = previous.map { "Previous summary: \($0)\n\nNew turns:\n\(transcript)" } ?? transcript
        return (try? await session.respond(to: prompt).content) ?? (previous ?? "")
    }
}
 
struct ChatTurn {
    enum Role: String { case user, assistant }
    let role: Role
    let text: String
}

import SwiftData
 
/// Layer 2 — chat history on disk: the full, uncompressed log, kept
/// independent from whatever the model's context window currently holds.
@Model
final class ChatMessage {
    var id: UUID = UUID()
    var role: ChatTurn.Role = .user
    var text: String = ""
    var createdAt: Date = .now
 
    init(role: ChatTurn.Role, text: String) {
        self.role = role
        self.text = text
    }
}

import SwiftData
 
/// Layer 3 — "Memory about you": a short, editable list of durable facts.
/// Lives on disk (and syncs via iCloud), independent of any single chat
/// session. The user can see, edit and delete every entry.
@Model
final class MemoryFact {
    var id: UUID = UUID()
    var text: String = ""
    var source: MemorySource = .manual
    var createdAt: Date = .now
 
    init(text: String, source: MemorySource = .manual) {
        self.text = text
        self.source = source
    }
}
 
enum MemorySource: String, Codable {
    case automatic  // extracted by the model from a conversation
    case manual     // added directly by the user
}
 
/// Best-effort auto-extraction: after a long-enough message, ask the model
/// for at most one durable fact. Never blocks the reply if it fails.
enum MemoryExtractor {
    static func extract(from text: String, existing: [MemoryFact]) async -> MemoryFact? {
        guard text.count > 25,
              case .available = SystemLanguageModel.default.availability else { return nil }
 
        let session = LanguageModelSession(instructions: """
        Extract ONE durable fact about the person from their message \
        (a preference, an important person, a habit worth remembering). \
        One short phrase, no quotes. If there is no durable fact, reply NONE.
        """)
        guard let response = try? await session.respond(to: text) else { return nil }
        let fact = response.content.trimmingCharacters(in: .whitespacesAndNewlines)
        guard !fact.isEmpty, fact.uppercased() != "NONE", fact.count < 120 else { return nil }
 
        let isDuplicate = existing.contains {
            $0.text.localizedCaseInsensitiveContains(fact)
        }
        return isDuplicate ? nil : MemoryFact(text: fact, source: .automatic)
    }
}

/// Packs all three layers into one instructions string that fits inside
/// the model's real context limit (Foundation Models: ~4096 tokens total,
/// input + output). Order matters: identity facts are cheap and high-value,
/// so they're never the first thing dropped.
struct PromptBudget {
    let totalTokens = 4096
    let reservedForResponse = 512
    let reservedForInstructions = 300
 
    var availableForContext: Int { totalTokens - reservedForResponse - reservedForInstructions }
 
    func assemble(memory: [MemoryFact], summary: String?, recentTurns: [ChatTurn]) -> String {
        var budget = availableForContext
        var blocks: [String] = []
 
        // Layer 3 — identity facts first: small, stable, high signal.
        if !memory.isEmpty {
            let text = memory.map { "- \($0.text)" }.joined(separator: "\n")
            blocks.append("About the person:\n\(text)")
            budget -= text.count / 4
        }
 
        // Layer 1, compressed part — the rolling summary of evicted turns.
        if let summary, !summary.isEmpty, budget > 200 {
            blocks.append("Earlier in the conversation: \(summary)")
            budget -= summary.count / 4
        }
 
        // Layer 1, verbatim tail — fill what's left, newest turns first.
        var tail: [String] = []
        for turn in recentTurns.reversed() {
            let cost = turn.text.count / 4
            guard cost < budget else { break }
            tail.insert("\(turn.role.rawValue): \(turn.text)", at: 0)
            budget -= cost
        }
        blocks.append(contentsOf: tail)
 
        return blocks.joined(separator: "\n\n")
    }
}