const { mockClient } = require('aws-sdk-client-mock');
const { BedrockRuntimeClient, ConverseCommand } = require('@aws-sdk/client-bedrock-runtime');
const { complete } = require('../src/llm');
const { ChapterSummarySchema } = require('../src/schemas');

const bedrock = mockClient(BedrockRuntimeClient);

function words(count) {
  return Array.from({ length: count }, (_, index) => `word${index}`).join(' ');
}

function summary(overrides = {}) {
  return {
    chapterNo: 1,
    title: 'Harbor',
    pageStart: 1,
    pageEnd: 3,
    oneLine: 'Ships leave the harbor.',
    summary: words(120),
    keyPoints: ['One', 'Two', 'Three'],
    entities: { people: ['Ada'], places: ['Harbor'], terms: ['Tide'] },
    questionsAnswered: ['Who leaves?', 'Where?', 'When?'],
    ...overrides,
  };
}

function toolResponse(input, name = 'chapter_summary') {
  return {
    output: { message: { content: [{ toolUse: { name, input } }] } },
    usage: { inputTokens: 11, outputTokens: 7 },
  };
}

describe('llm repair', () => {
  let client;

  beforeEach(() => {
    bedrock.reset();
    client = new BedrockRuntimeClient({ region: 'ap-south-1' });
  });

  it('retries once when the first tool input fails validation', async () => {
    bedrock
      .on(ConverseCommand)
      .resolvesOnce(toolResponse({ ...summary(), summary: 'too short' }))
      .resolvesOnce(toolResponse(summary()));
    const result = await complete({
      client,
      modelId: 'test-model',
      system: 'system',
      user: '<book-text>\nHello\n</book-text>',
      schema: ChapterSummarySchema,
      toolName: 'chapter_summary',
      jsonSchema: { type: 'object' },
    });
    expect(result.data.summary.split(/\s+/)).toHaveLength(120);
    expect(result.usage.inputTokens).toBe(22);
    expect(bedrock.commandCalls(ConverseCommand)).toHaveLength(2);
  });

  it('throws after the repair attempt is still invalid', async () => {
    bedrock.on(ConverseCommand).resolves(toolResponse({ ...summary(), summary: 'still short' }));
    await expect(
      complete({
        client,
        modelId: 'test-model',
        system: 'system',
        user: 'text',
        schema: ChapterSummarySchema,
        toolName: 'chapter_summary',
        jsonSchema: { type: 'object' },
      }),
    ).rejects.toThrow(/after repair/);
    expect(bedrock.commandCalls(ConverseCommand)).toHaveLength(2);
  });
});