/**
 * Token Counting Utilities
 * 
 * Implements a heuristic-based token counter (approx 4 chars per token)
 * to avoid shipping a heavy tokenizer like Tiktoken to the client.
 * 
 * Risk R-2.2 Mitigation: Prevents Context Window Overflow.
 */

export const ESTIMATED_CHARS_PER_TOKEN = 4;

/**
 * Estimates the number of tokens in a string.
 * Uses the industry standard approximation of 4 characters per token for English text.
 */
export function estimateTokenCount(text: string): number {
    if (!text) return 0;
    return Math.ceil(text.length / ESTIMATED_CHARS_PER_TOKEN);
}

export interface ChatMessage {
    role: 'user' | 'system' | 'assistant';
    content: string;
}

/**
 * Estimates the token count for a list of messages.
 * Adds overhead for message formatting (approx 4 tokens per message).
 */
export function countChatTokens(messages: ChatMessage[]): number {
    let total = 0;
    for (const msg of messages) {
        // Content tokens
        total += estimateTokenCount(msg.content);
        // Overhead (role + structural tokens)
        total += 4;
    }
    return total;
}

/**
 * Truncates a chat history to fit within a specific token limit.
 * Keeps the system prompt (first message) and the most recent messages.
 * Drops the oldest messages from the middle.
 */
export function truncateChatHistory(
    messages: ChatMessage[],
    maxTokens: number,
    systemPromptTokens: number = 0
): ChatMessage[] {
    // 1. Calculate total current tokens
    const currentTokens = countChatTokens(messages);

    // If we are within limits, return original
    if (currentTokens <= maxTokens) {
        return messages;
    }

    // 2. Identify critical messages (System Prompt + Last User Message)
    const result: ChatMessage[] = [];
    let reservedTokens = 0;

    // Always keep the first message if it is a system prompt
    if (messages.length > 0 && messages[0].role === 'system') {
        result.push(messages[0]);
        reservedTokens += estimateTokenCount(messages[0].content) + 4;
    }

    // Ensure we have budget
    const budget = maxTokens - reservedTokens - systemPromptTokens;
    if (budget <= 0) {
        console.warn('Token budget too small for even the system prompt.');
        return result;
    }

    // 3. Add messages from the end backwards until budget is full
    // We skip the first message if we already added it (as system prompt)
    const startIndex = (result.length > 0) ? 1 : 0;
    const messagesToConsider = messages.slice(startIndex).reverse();

    const recentMessages: ChatMessage[] = [];
    let used = 0;

    for (const msg of messagesToConsider) {
        const cost = estimateTokenCount(msg.content) + 4;
        if (used + cost <= budget) {
            recentMessages.unshift(msg);
            used += cost;
        } else {
            break; // Stop adding once we hit the limit
        }
    }

    // 4. Combine
    return [...result, ...recentMessages];
}
