const { mockClient } = require('aws-sdk-client-mock');
const { BedrockRuntimeClient, ConverseStreamCommand } = require('@aws-sdk/client-bedrock-runtime');
const { createEngine } = require('../src/engine');
const { historyMessages, prefixBeforeCheckpoint } = require('../src/prompt/build');
const { MemorySessionStore } = require('../src/sessions');
const { BookCache } = require('../src/bookCache');
const { sentenceStream } = require('../src/sentenceStream');
const { chatConfig, memoryStorage, textStream, usage } = require('./fixture');

const bedrock = mockClient(BedrockRuntimeClient);

async function collect(iterable) {
  const events = [];
  for await (const event of iterable) events.push(event);
  return events;
}

describe('engine', () => {
  let client;

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

  it('persists an interrupted turn', async () => {
    const abort = new AbortController();
    bedrock.on(ConverseStreamCommand).callsFake(async () => ({
      stream: (async function* () {
        yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: 'Hello' } } };
        abort.abort();
        yield { contentBlockDelta: { contentBlockIndex: 0, delta: { text: ' more' } } };
        yield { metadata: { usage: usage({ outputTokens: 2 }) } };
      })(),
    }));
    const sessions = new MemorySessionStore();
    const storage = memoryStorage('3');
    const engine = createEngine({
      config: chatConfig(),
      logger: { info() {}, warn() {} },
      storage,
      books: { async get() { return { bookId: 3, tokenCount: 10 }; } },
      sessions,
      bedrock: client,
      opensearch: null,
      bookCache: new BookCache(storage),
    });
    const session = await engine.startSession({ bookId: 3, channel: 'chat', modeOverride: 'full' });
    const events = await collect(engine.answer({ sessionId: session.sessionId, userText: 'Hi', signal: abort.signal }));
    expect(events.find((event) => event.type === 'final').interrupted).toBe(true);
    const saved = await sessions.get(session.sessionId);
    expect(saved.history[1].text).toBe('Hello');
    expect(saved.history[1].interrupted).toBe(true);
    expect(saved.history[1].rawText).toBeUndefined();
  });

  it('sends markers back on the next turn and leaves the cached prefix unchanged', async () => {
    const inputs = [];
    bedrock.on(ConverseStreamCommand).callsFake(async (input) => {
      inputs.push(input);
      return {
        stream: textStream('It is listed [[p:1]].', usage({ outputTokens: 6 })),
      };
    });
    const sessions = new MemorySessionStore();
    const storage = memoryStorage('3');
    const engine = createEngine({
      config: chatConfig(),
      logger: { info() {}, warn() {} },
      storage,
      books: { async get() { return { bookId: 3, tokenCount: 10, status: 'completed' }; } },
      sessions,
      bedrock: client,
      opensearch: null,
      bookCache: new BookCache(storage),
    });
    const session = await engine.startSession({ bookId: 3, channel: 'chat', modeOverride: 'full' });
    const first = await collect(engine.answer({ sessionId: session.sessionId, userText: 'Where?' }));
    expect(first.filter((event) => event.type === 'text_delta').map((event) => event.text).join('')).not.toContain('[[');
    expect(first.find((event) => event.type === 'final').rawText).toBeUndefined();
    const saved = await sessions.get(session.sessionId);
    expect(saved.history[1].rawText).toContain('[[p:1]]');
    expect(saved.history[1].text).toBe('It is listed.');
    const replay = historyMessages({ history: saved.history });
    expect(replay[1].content[0].text).toContain('[[p:1]]');
    await collect(engine.answer({ sessionId: session.sessionId, userText: 'And the name?' }));
    const prefix = (input) => JSON.stringify(prefixBeforeCheckpoint(input.system));
    expect(prefix(inputs[0])).toBe(prefix(inputs[1]));
    const secondAssistant = inputs[1].messages.find((message) => message.role === 'assistant');
    expect(secondAssistant.content[0].text).toContain('[[p:1]]');
  });
});

describe('sentenceStream', () => {
  it('keeps abbreviations and decimals intact and flushes the tail', async () => {
    async function* deltas() {
      yield { type: 'text_delta', text: 'Dr. Ada left. ' };
      yield { type: 'text_delta', text: 'The value is 3.14 today. Next' };
    }
    const sentences = [];
    for await (const event of sentenceStream(deltas())) {
      if (event.type === 'sentence') sentences.push(event.text);
    }
    expect(sentences).toEqual(['Dr. Ada left.', 'The value is 3.14 today.', 'Next']);
  });
});