//
//  ChatSession.swift
//  Zyquo Local
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import MLXLMCommon

/// Zyquo Local's multi-turn session: converts persisted conversation history
/// into the model's chat-template form, truncates oldest turns when the
/// context window would overflow (always keeping the system prompt), and
/// reuses the KV cache across turns via the underlying MLX session.
///
/// Not thread-safe by design — owned and accessed only by `InferenceEngine`.
final class ChatSession {
    /// Conversation this session is bound to.
    let conversationID: UUID
    private let contextWindow: Int
    private let mlxSession: MLXLMCommon.ChatSession

    /// Rough chars-per-token estimate used for truncation budgeting.
    private static let charsPerToken = 3.5
    /// Fraction of the context window budgeted for history (rest is reserved
    /// for the generated response).
    private static let historyBudgetFraction = 0.7

    /// Number of tokens estimated to be used by history + system prompt.
    private(set) var estimatedHistoryTokens: Int

    init(
        container: ModelContainer,
        conversation: Conversation,
        contextWindow: Int,
        params: GenerationParams
    ) {
        self.conversationID = conversation.id
        self.contextWindow = contextWindow

        let truncated = Self.truncatedHistory(
            messages: conversation.messages,
            contextWindow: contextWindow
        )
        self.estimatedHistoryTokens = Self.estimateTokens(
            truncated.map(\.content).joined() + (conversation.systemPrompt ?? "")
        )

        let history: [Chat.Message] = truncated.map { message in
            switch message.role {
            case .system: .system(message.content)
            case .user: .user(message.content)
            case .assistant: .assistant(message.content)
            }
        }
        self.mlxSession = MLXLMCommon.ChatSession(
            container,
            instructions: conversation.systemPrompt,
            history: history,
            generateParameters: params.toMLX()
        )
    }

    /// Streams a response to a new user prompt, honoring per-call parameters
    /// without losing the KV cache.
    func stream(prompt: String, params: GenerationParams) -> AsyncThrowingStream<Generation, Error> {
        mlxSession.generateParameters = params.toMLX()
        estimatedHistoryTokens += Self.estimateTokens(prompt)
        return mlxSession.streamDetails(to: prompt)
    }

    /// Estimated context usage in tokens (for the UI context bar).
    var estimatedContextUsage: (used: Int, window: Int) {
        (estimatedHistoryTokens, contextWindow)
    }

    func noteResponse(_ text: String) {
        estimatedHistoryTokens += Self.estimateTokens(text)
    }

    // MARK: - Truncation

    static func estimateTokens(_ text: String) -> Int {
        Int(Double(text.count) / charsPerToken) + 1
    }

    /// Drops oldest non-system turns until the estimated history fits the
    /// budgeted share of the context window. The system prompt (a leading
    /// system message, if any) is always kept.
    static func truncatedHistory(messages: [Message], contextWindow: Int) -> [Message] {
        let budget = Int(Double(contextWindow) * historyBudgetFraction)
        var system: [Message] = []
        var turns: [Message] = []
        for m in messages {
            if m.role == .system && turns.isEmpty {
                system.append(m)
            } else {
                turns.append(m)
            }
        }
        func total(_ list: [Message]) -> Int {
            list.reduce(0) { $0 + estimateTokens($1.content) }
        }
        while turns.count > 1, total(system) + total(turns) > budget {
            turns.removeFirst()
        }
        return system + turns
    }
}