const { mockClient } = require('aws-sdk-client-mock');
const { BedrockRuntimeClient, CountTokensCommand } = require('@aws-sdk/client-bedrock-runtime');
const { countBook, estimateTokens } = require('../src/tokens');

const bedrock = mockClient(BedrockRuntimeClient);

describe('token counts', () => {
  beforeEach(() => {
    bedrock.reset();
  });

  it('falls back to chars/4 when CountTokens fails', async () => {
    bedrock.on(CountTokensCommand).rejects(new Error('model does not support CountTokens'));
    const client = new BedrockRuntimeClient({ region: 'ap-south-1' });
    const marked = '[Page 1]\nHello world\n';
    const result = await countBook({
      client,
      modelId: 'test-count',
      markedText: marked,
      chapters: [{ chapterNo: 1, startPage: 1, endPage: 1 }],
      pages: [{ page: 1, text: 'Hello world' }],
    });
    expect(result.tokenCountMethod).toBe('estimate');
    expect(result.tokenCount).toBe(estimateTokens(marked));
    expect(result.chapters[0].tokenCountMethod).toBe('estimate');
  });

  it('uses Bedrock inputTokens when the call succeeds', async () => {
    bedrock.on(CountTokensCommand).resolves({ inputTokens: 42 });
    const client = new BedrockRuntimeClient({ region: 'ap-south-1' });
    const result = await countBook({
      client,
      modelId: 'test-count',
      markedText: 'Hello',
      chapters: [{ chapterNo: 1, startPage: 1, endPage: 1 }],
      pages: [{ page: 1, text: 'Hello' }],
    });
    expect(result.tokenCountMethod).toBe('bedrock');
    expect(result.tokenCount).toBe(42);
    expect(result.chapters[0].tokenCount).toBe(42);
  });
});