import { describe, expect, it } from 'vitest';
import type { Message, ToolDef } from '../../llm/openai-compat.js';
import {
  estimateMessagesTokens,
  estimateMessageTokens,
  estimatePromptTokens,
  estimateTokensFromChars,
  estimateTokensFromText,
  estimateToolsTokens,
  IMAGE_CONTENT_TOKENS,
  UNKNOWN_CHARS_TO_TOKENS_FACTOR,
} from './token-estimate.js';

describe('estimateTokensFromChars', () => {
  it('returns 0 for 0 chars', () => {
    expect(estimateTokensFromChars(0)).toBe(0);
  });

  it('applies the unknown-content factor with ceiling', () => {
    expect(estimateTokensFromChars(1)).toBe(Math.ceil(1 * UNKNOWN_CHARS_TO_TOKENS_FACTOR)); // 2
    expect(estimateTokensFromChars(5)).toBe(6);   // 5 * 1.2 = 6 exactly
    expect(estimateTokensFromChars(10)).toBe(12); // 10 * 1.2 = 12 exactly
    expect(estimateTokensFromChars(11)).toBe(14); // 13.2 -> ceil 14
  });
});

describe('estimateTokensFromText', () => {
  it('returns 0 for empty string', () => {
    expect(estimateTokensFromText('')).toBe(0);
  });

  it('counts ASCII at ~3.5 chars per token', () => {
    expect(estimateTokensFromText('a'.repeat(7))).toBe(2);   // 7 / 3.5 = 2 exactly
    expect(estimateTokensFromText('a'.repeat(35))).toBe(10); // 35 / 3.5 = 10 exactly
    expect(estimateTokensFromText('a')).toBe(1);             // 0.286 -> ceil 1
  });

  it('counts CJK at 1.2 tokens per char', () => {
    expect(estimateTokensFromText('あ')).toBe(2);    // 1.2 -> ceil 2
    expect(estimateTokensFromText('日本語')).toBe(4); // 3.6 -> ceil 4
    expect(estimateTokensFromText('カタカナ')).toBe(5); // 4.8 -> ceil 5
  });

  it('classifies full-width forms as CJK', () => {
    // 'Ａ' (U+FF21) and '！' (U+FF01) are in the full-width block.
    expect(estimateTokensFromText('ＡＢ')).toBe(3); // 2.4 -> ceil 3
  });

  it('counts other non-ASCII (accented latin, emoji) at 1 token per char', () => {
    expect(estimateTokensFromText('é')).toBe(1);
    expect(estimateTokensFromText('éà')).toBe(2);
  });

  it('iterates astral code points once (emoji is one char, not two surrogates)', () => {
    expect(estimateTokensFromText('😀')).toBe(1);
    expect(estimateTokensFromText('😀😀😀')).toBe(3);
  });

  it('mixes script classes additively before a single ceil', () => {
    // 'abc' (3 ascii -> 0.857) + 'あ' (1.2) = 2.057 -> ceil 3
    expect(estimateTokensFromText('abcあ')).toBe(3);
  });

  it('estimates Japanese text far higher per char than ASCII text', () => {
    const ascii = estimateTokensFromText('a'.repeat(100));
    const cjk = estimateTokensFromText('あ'.repeat(100));
    expect(ascii).toBe(29);  // 100 / 3.5 -> ceil
    expect(cjk).toBe(120);   // 100 * 1.2
  });
});

describe('estimateMessageTokens', () => {
  it('charges role text plus a fixed 8-token overhead for an empty message', () => {
    // 'user' = 4 ascii chars -> ceil(4/3.5) = 2, plus 8 overhead.
    expect(estimateMessageTokens({ role: 'user' })).toBe(10);
  });

  it('adds string content tokens', () => {
    const base = estimateMessageTokens({ role: 'user' });
    const withContent = estimateMessageTokens({ role: 'user', content: 'a'.repeat(35) });
    expect(withContent).toBe(base + 10);
  });

  it('sums text parts in array content', () => {
    const msg: Message = {
      role: 'user',
      content: [
        { type: 'text', text: 'a'.repeat(7) },
        { type: 'text', text: 'b'.repeat(7) },
      ],
    };
    expect(estimateMessageTokens(msg)).toBe(estimateMessageTokens({ role: 'user' }) + 4);
  });

  it('charges a flat image budget per image part', () => {
    const msg: Message = {
      role: 'user',
      content: [{ type: 'image_url', image_url: { url: 'data:image/png;base64,xxxx' } }],
    };
    expect(estimateMessageTokens(msg)).toBe(estimateMessageTokens({ role: 'user' }) + IMAGE_CONTENT_TOKENS);
  });

  it('counts tool_call_id and name on tool messages', () => {
    const base = estimateMessageTokens({ role: 'tool' });
    const msg: Message = {
      role: 'tool',
      tool_call_id: 'c'.repeat(7),  // 2 tokens
      name: 'n'.repeat(7),          // 2 tokens
    };
    expect(estimateMessageTokens(msg)).toBe(base + 4);
  });

  it('counts tool_calls id, function name, and arguments', () => {
    const base = estimateMessageTokens({ role: 'assistant' });
    const msg: Message = {
      role: 'assistant',
      tool_calls: [{
        id: 'i'.repeat(7),                                 // 2 tokens
        type: 'function',
        function: { name: 'f'.repeat(7), arguments: 'a'.repeat(35) }, // 2 + 10 tokens
      }],
    };
    expect(estimateMessageTokens(msg)).toBe(base + 14);
  });

  it('handles a message combining content and multiple tool calls', () => {
    const base = estimateMessageTokens({ role: 'assistant' });
    const call = {
      id: 'i'.repeat(7),
      type: 'function' as const,
      function: { name: 'f'.repeat(7), arguments: 'a'.repeat(7) },
    };
    const msg: Message = {
      role: 'assistant',
      content: 'a'.repeat(7),
      tool_calls: [call, call],
    };
    // content 2 + 2 calls * (2 + 2 + 2)
    expect(estimateMessageTokens(msg)).toBe(base + 2 + 12);
  });
});

describe('estimateMessagesTokens', () => {
  it('returns 0 for an empty list', () => {
    expect(estimateMessagesTokens([])).toBe(0);
  });

  it('sums per-message estimates', () => {
    const m1: Message = { role: 'user', content: 'hello' };
    const m2: Message = { role: 'assistant', content: 'world' };
    expect(estimateMessagesTokens([m1, m2]))
      .toBe(estimateMessageTokens(m1) + estimateMessageTokens(m2));
  });
});

describe('estimateToolsTokens', () => {
  it('estimates from the JSON serialization of the tool list', () => {
    const tools: ToolDef[] = [{
      type: 'function',
      function: { name: 'Read', description: 'Read a file.', parameters: { type: 'object' } },
    }];
    expect(estimateToolsTokens(tools)).toBe(estimateTokensFromText(JSON.stringify(tools)));
    expect(estimateToolsTokens(tools)).toBeGreaterThan(0);
  });

  it('costs a small but non-zero amount for an empty tool list ("[]" is 2 chars)', () => {
    expect(estimateToolsTokens([])).toBe(1);
  });
});

describe('estimatePromptTokens', () => {
  it('is the sum of message and tool estimates', () => {
    const messages: Message[] = [{ role: 'system', content: 'You are helpful.' }];
    const tools: ToolDef[] = [{
      type: 'function',
      function: { name: 'X', description: 'd', parameters: {} },
    }];
    expect(estimatePromptTokens(messages, tools))
      .toBe(estimateMessagesTokens(messages) + estimateToolsTokens(tools));
  });

  it('handles empty messages and tools', () => {
    expect(estimatePromptTokens([], [])).toBe(1); // only "[]" tools serialization
  });
});
