import { estimateTokenCount, countChatTokens, truncateChatHistory, ChatMessage } from './token-utils';
import { describe, it, expect } from 'vitest';

describe('Token Utilities', () => {
    describe('estimateTokenCount', () => {
        it('should return 0 for empty string', () => {
            expect(estimateTokenCount('')).toBe(0);
        });

        it('should estimate 4 chars as 1 token', () => {
            expect(estimateTokenCount('word')).toBe(1);
        });

        it('should round up', () => {
            expect(estimateTokenCount('abcde')).toBe(2); // 5 chars -> 1.25 -> 2 tokens
        });
    });

    describe('countChatTokens', () => {
        it('should count content plus overhead', () => {
            const messages: ChatMessage[] = [
                { role: 'user', content: 'hello' } // 5 chars (2 tokens) + 4 overhead = 6
            ];
            expect(countChatTokens(messages)).toBe(6);
        });
    });

    describe('truncateChatHistory', () => {
        const systemMsg: ChatMessage = { role: 'system', content: 'Act like a teacher' };
        const userMsg1: ChatMessage = { role: 'user', content: 'Hello' };
        const assistantMsg: ChatMessage = { role: 'assistant', content: 'Hi there' };
        const userMsg2: ChatMessage = { role: 'user', content: 'How are you?' };

        it('should return all messages if within limit', () => {
            const history = [systemMsg, userMsg1, assistantMsg, userMsg2];
            // Total tokens should be roughly:
            // sys: 18 chars -> 5 tokens + 4 = 9
            // u1: 5 chars -> 2 tokens + 4 = 6
            // a1: 8 chars -> 2 tokens + 4 = 6
            // u2: 12 chars -> 3 tokens + 4 = 7
            // Total: 28 tokens
            const result = truncateChatHistory(history, 100);
            expect(result).toHaveLength(4);
        });

        it('should always keep the system prompt', () => {
            const history = [systemMsg, userMsg1, assistantMsg, userMsg2];
            // Limit to 15. System is 9. Budget left 6.
            // u2 is 7 -> too big?
            // u1 is 6 -> fits?
            // Wait, we reverse: most recent first.
            // userMsg2 is newest. Cost 7. Limit 15 - 9 (sys) = 6 budget.
            // userMsg2 needs 7. 7 > 6. So it should be skipped.
            // Wait, logic says "stop adding".

            const result = truncateChatHistory(history, 15);
            expect(result).toHaveLength(1);
            expect(result[0]).toEqual(systemMsg);
        });

        it('should prioritize recent messages', () => {
            const history = [systemMsg, userMsg1, assistantMsg, userMsg2];
            // Limit 20 (Sys=9). Budget 11.
            // u2 (7 tokens) -> Budget 4. Added.
            // a1 (6 tokens) -> Budget -2. Skipped.

            const result = truncateChatHistory(history, 20);
            expect(result).toHaveLength(2);
            expect(result[0]).toEqual(systemMsg);
            expect(result[1]).toEqual(userMsg2);
        });
    });
});
