import { describe, expect, it, vi } from 'vitest';

import {
  ABORT_CHUNK,
  convertIterableToStream,
  createCallbacksTransformer,
  createFirstErrorHandleTransformer,
  createSSEDataExtractor,
  createSSEProtocolTransformer,
  createTokenSpeedCalculator,
  FIRST_CHUNK_ERROR_KEY,
  readableFromAsyncIterable,
} from './protocol';

describe('createSSEDataExtractor', () => {
  // Helper function to convert string to Uint8Array
  const stringToUint8Array = (str: string): Uint8Array => {
    return new TextEncoder().encode(str);
  };

  // Helper function to process chunks through transformer
  const processChunk = async (transformer: TransformStream, chunk: Uint8Array) => {
    const results: any[] = [];
    const readable = new ReadableStream({
      start(controller) {
        controller.enqueue(chunk);
        controller.close();
      },
    });

    const writable = new WritableStream({
      write(chunk) {
        results.push(chunk);
      },
    });

    await readable.pipeThrough(transformer).pipeTo(writable);

    return results;
  };

  it('should correctly transform single SSE data line', async () => {
    const transformer = createSSEDataExtractor();
    const input = 'data: {"message": "hello"}\n';
    const chunk = stringToUint8Array(input);

    const results = await processChunk(transformer, chunk);

    expect(results).toEqual([{ message: 'hello' }]);
  });

  it('should handle multiple SSE data lines', async () => {
    const transformer = createSSEDataExtractor();
    const input = `data: {"message": "hello"}\ndata: {"message": "world"}\n`;
    const chunk = stringToUint8Array(input);

    const results = await processChunk(transformer, chunk);

    expect(results).toEqual([{ message: 'hello' }, { message: 'world' }]);
  });

  it('should ignore non-data lines', async () => {
    const transformer = createSSEDataExtractor();
    const input = `id: 1\ndata: {"message": "hello"}\nevent: message\n`;
    const chunk = stringToUint8Array(input);

    const results = await processChunk(transformer, chunk);

    expect(results).toEqual([{ message: 'hello' }]);
  });

  it('should skip [DONE] heartbeat messages', async () => {
    const transformer = createSSEDataExtractor();
    const input = `data: {"message": "hello"}\ndata: [DONE]\ndata: {"message": "world"}\n`;
    const chunk = stringToUint8Array(input);

    const results = await processChunk(transformer, chunk);

    expect(results).toEqual([{ message: 'hello' }, { message: 'world' }]);
  });

  it('should handle invalid JSON gracefully', async () => {
    const transformer = createSSEDataExtractor();
    const input = `data: {"message": "hello"}\ndata: invalid-json\ndata: {"message": "world"}\n`;
    const chunk = stringToUint8Array(input);

    const results = await processChunk(transformer, chunk);

    expect(results).toEqual([{ message: 'hello' }, { message: 'world' }]);
  });

  it('should handle empty data lines', async () => {
    const transformer = createSSEDataExtractor();
    const input = `data: \ndata: {"message": "hello"}\ndata: \n`;
    const chunk = stringToUint8Array(input);

    const results = await processChunk(transformer, chunk);

    expect(results).toEqual([{ message: 'hello' }]);
  });

  it('should process large chunks of data correctly', async () => {
    const transformer = createSSEDataExtractor();
    const messages = Array.from({ length: 100 })
      .fill(null)
      .map((_, i) => `data: {"message": "message${i}"}\n`)
      .join('');
    const chunk = stringToUint8Array(messages);

    const results = await processChunk(transformer, chunk);

    expect(results).toHaveLength(100);
    expect(results[0]).toEqual({ message: 'message0' });
    expect(results[99]).toEqual({ message: 'message99' });
  });

  describe('real world data', () => {
    it('should convert azure ai data', async () => {
      const chunks = [
        `data: {"choices":[{"delta":{"content":"","reasoning_content":null,"role":"assistant","tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714651,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"\u003Cthink\u003E","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714651,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"\n\n","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714651,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"\u003C/think\u003E","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714651,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"\n\n","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714651,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"Hello","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714651,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"!","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" How","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" can","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" I","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" assist","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" you","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" today","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"?","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":" ","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"😊","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":null,"index":0,"logprobs":null,"matched_stop":null}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[{"delta":{"content":"","reasoning_content":null,"role":null,"tool_calls":null},"finish_reason":"stop","index":0,"logprobs":null,"matched_stop":1}],"created":1739714652,"id":"1392a93d52c3483ea872d0ab2aaff7d7","model":"DeepSeek-R1","object":"chat.completion.chunk","usage":null}\n`,
        `data: {"choices":[],"id":"79fca0de792a4ffb8ec836442a2a42c0","model":"DeepSeek-R1","usage":{"completion_tokens":16,"prompt_tokens":4,"total_tokens":20}}\n`,
        `data: [DONE]`,
      ];

      const transformer = createSSEDataExtractor();

      const results = await processChunk(transformer, stringToUint8Array(chunks.join('')));
      expect(results).matchSnapshot();
    });
  });
});

describe('createTokenSpeedCalculator', async () => {
  // Mock the param from caller - 1000 to avoid div 0
  const inputStartAt = Date.now() - 1000;

  // Helper function to process chunks through transformer
  const processChunk = async (transformer: TransformStream, chunk: any) => {
    const results: any[] = [];
    const readable = new ReadableStream({
      start(controller) {
        controller.enqueue(chunk);
        controller.close();
      },
    });

    const writable = new WritableStream({
      write(chunk) {
        results.push(chunk);
      },
    });

    await readable.pipeThrough(transformer).pipeTo(writable);

    return results;
  };

  it('should calculate token speed correctly', async () => {
    const chunks = [
      { data: '', id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy', type: 'text' },
      { data: 'hi', id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy', type: 'text' },
      { data: 'stop', id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy', type: 'stop' },
      {
        data: {
          inputTextTokens: 9,
          outputTextTokens: 1,
          totalInputTokens: 9,
          totalOutputTokens: 1,
          totalTokens: 10,
        },
        id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy',
        type: 'usage',
      },
    ];

    const transformer = createTokenSpeedCalculator((v) => v, { inputStartAt });
    const results = await processChunk(transformer, chunks);
    expect(results).toHaveLength(chunks.length + 1);
    const speedChunk = results.at(-1);
    expect(speedChunk.id).toBe('output_speed');
    expect(speedChunk.type).toBe('speed');
    expect(speedChunk.data.tps).not.toBeNaN();
    expect(speedChunk.data.ttft).not.toBeNaN();
  });

  it('should not calculate token speed if no usage', async () => {
    const chunks = [
      { data: '', id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy', type: 'text' },
      { data: 'hi', id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy', type: 'text' },
      { data: 'stop', id: 'chatcmpl-BKO1bogylHvMaYfETjTAzrCguYwZy', type: 'stop' },
    ];

    const transformer = createTokenSpeedCalculator((v) => v, { inputStartAt });
    const results = await processChunk(transformer, chunks);
    expect(results).toHaveLength(chunks.length);
  });

  it('should calculate token speed considering outputImageTokens when totalOutputTokens is missing', async () => {
    const chunks = [
      { data: '', id: 'chatcmpl-image-1', type: 'text' },
      { data: 'hi', id: 'chatcmpl-image-1', type: 'text' },
      { data: 'stop', id: 'chatcmpl-image-1', type: 'stop' },
      {
        data: {
          inputTextTokens: 9,
          outputTextTokens: 1,
          outputImageTokens: 4,
          totalInputTokens: 9,
          // totalOutputTokens intentionally omitted to force summation path
          totalTokens: 13,
        },
        id: 'chatcmpl-image-1',
        type: 'usage',
      },
    ];

    const transformer = createTokenSpeedCalculator((v) => v, { inputStartAt });
    const results = await processChunk(transformer, chunks);

    // should push an extra speed chunk
    expect(results).toHaveLength(chunks.length + 1);
    const speedChunk = results.at(-1);
    expect(speedChunk.id).toBe('output_speed');
    expect(speedChunk.type).toBe('speed');
    // tps and ttft should be numeric (avoid flakiness if interval is 0ms)
    expect(speedChunk.data.tps).not.toBeNaN();
    expect(speedChunk.data.ttft).not.toBeNaN();
  });
});

describe('convertIterableToStream', () => {
  const drain = async (readable: ReadableStream<any>) => {
    const reader = readable.getReader();
    const chunks: any[] = [];
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }
    return chunks;
  };

  it('should surface errors from subsequent pulls as error chunks', async () => {
    async function* erroringStream() {
      yield 'first';
      throw new Error('rate limit');
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream()).pipeThrough(createFirstErrorHandleTransformer()),
    );

    expect(chunks[0]).toBe('first');
    expect(chunks[1][FIRST_CHUNK_ERROR_KEY]).toBe(true);
    expect(chunks[1].message).toBe('rate limit');
  });

  it('should enrich error chunks with provider/model context', async () => {
    async function* erroringStream() {
      yield 'first';
      throw new Error('connection reset');
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream(), {
        model: 'deepseek-v4-flash',
        provider: 'lobehub',
      }).pipeThrough(createFirstErrorHandleTransformer(undefined, 'lobehub')),
    );

    expect(chunks[1].message).toBe('connection reset');
    expect(chunks[1].provider).toBe('lobehub');
    expect(chunks[1].model).toBe('deepseek-v4-flash');
  });

  it('should extract parse position from JSON SyntaxError messages', async () => {
    async function* erroringStream() {
      yield 'first';
      // Reproduce the V8 JSON.parse SyntaxError shape that surfaces from the
      // OpenAI SDK iterator when an upstream SSE chunk contains an illegal
      // backslash escape — see LobeHub op_1778403331540 for a real instance.
      throw new SyntaxError(
        'Bad escaped character in JSON at position 160050 (line 1 column 160051)',
      );
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream(), { provider: 'lobehub' }).pipeThrough(
        createFirstErrorHandleTransformer(undefined, 'lobehub'),
      ),
    );

    expect(chunks[1].name).toBe('SyntaxError');
    expect(chunks[1].parsePosition).toBe(160050);
  });

  it('should surface error.cause when present', async () => {
    async function* erroringStream() {
      yield 'first';
      throw new Error('wrapper', { cause: new SyntaxError('inner parse failure') });
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream()).pipeThrough(createFirstErrorHandleTransformer()),
    );

    expect(chunks[1].causeName).toBe('SyntaxError');
    expect(chunks[1].causeMessage).toBe('inner parse failure');
  });

  it('should extract parsePosition from a wrapped SyntaxError cause', async () => {
    // Many provider SDKs rethrow JSON.parse failures wrapped in their own
    // error class (e.g. APIError) — the outer name is no longer
    // 'SyntaxError', so the offset has to be pulled from `cause`.
    class APIError extends Error {
      constructor(message: string, options?: { cause?: unknown }) {
        super(message, options);
        this.name = 'APIError';
      }
    }

    async function* erroringStream() {
      yield 'first';
      throw new APIError('upstream failed', {
        cause: new SyntaxError(
          'Bad escaped character in JSON at position 160050 (line 1 column 160051)',
        ),
      });
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream()).pipeThrough(createFirstErrorHandleTransformer()),
    );

    expect(chunks[1].name).toBe('APIError');
    expect(chunks[1].causeName).toBe('SyntaxError');
    expect(chunks[1].parsePosition).toBe(160050);
  });

  it('should not throw when cause contains BigInt or circular refs', async () => {
    // structuredClone accepts both of these; JSON.stringify does not. If the
    // outer stringify in buildStreamErrorPayload fails, the FIRST_CHUNK_ERROR
    // chunk is never emitted and the stream silently dies — test that the
    // diagnostic path stays intact.
    const circular: Record<string, unknown> = { kind: 'detail' };
    circular.self = circular;
    const badCause = { big: 9_007_199_254_740_993n, ref: circular };

    async function* erroringStream() {
      yield 'first';
      throw new Error('upstream blew up', { cause: badCause });
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream()).pipeThrough(createFirstErrorHandleTransformer()),
    );

    expect(chunks[1].message).toBe('upstream blew up');
    // Cause is an object, not an Error, so it goes through `toJsonSafe`.
    expect(chunks[1].cause).toBeDefined();
    expect(chunks[1].cause.big).toBe('9007199254740993');
    expect(chunks[1].cause.ref.self).toBe('[Circular]');
  });

  it('should emit ABORT_CHUNK when AbortError occurs during pull', async () => {
    async function* abortingStream() {
      yield 'first';
      const err = new Error('The operation was aborted');
      err.name = 'AbortError';
      throw err;
    }

    const readable = convertIterableToStream(abortingStream());
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks).toEqual(['first', ABORT_CHUNK]);
  });

  it('should emit ABORT_CHUNK when "Request was aborted" error occurs during pull', async () => {
    async function* abortingStream() {
      yield 'first';
      throw new Error('Request was aborted.');
    }

    const readable = convertIterableToStream(abortingStream());
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks).toEqual(['first', ABORT_CHUNK]);
  });

  it('should emit ABORT_CHUNK when abort error occurs during start', async () => {
    async function* abortingStream(): AsyncGenerator<string> {
      yield* []; // eslint: require-yield
      throw new Error('Request was aborted.');
    }

    const readable = convertIterableToStream(abortingStream());
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks).toEqual([ABORT_CHUNK]);
  });

  it('should emit ABORT_CHUNK when "cancelled" error occurs during pull', async () => {
    async function* cancelledStream() {
      yield 'data';
      throw new Error('The request was cancelled');
    }

    const readable = convertIterableToStream(cancelledStream());
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks).toEqual(['data', ABORT_CHUNK]);
  });

  it('should emit ABORT_CHUNK for AbortError-named throw without a message', async () => {
    async function* abortingStream() {
      yield 'first';
      // some SDKs throw bare objects rather than Error instances
      throw { name: 'AbortError' };
    }

    const readable = convertIterableToStream(abortingStream());
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks).toEqual(['first', ABORT_CHUNK]);
  });

  it('should fall back to error chunk when iterator throws a non-Error value', async () => {
    async function* erroringStream() {
      yield 'first';
      // a string throw must not crash the abort check — it should still
      // surface as a serialized first-chunk error instead of killing the pipe
      throw 'upstream exploded';
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream()).pipeThrough(createFirstErrorHandleTransformer()),
    );

    expect(chunks[0]).toBe('first');
    expect(chunks[1][FIRST_CHUNK_ERROR_KEY]).toBe(true);
  });

  it('should fall back to error chunk when thrown object has no message', async () => {
    async function* erroringStream() {
      yield 'first';
      throw { code: 'ECONNRESET' };
    }

    const chunks = await drain(
      convertIterableToStream(erroringStream()).pipeThrough(createFirstErrorHandleTransformer()),
    );

    expect(chunks[0]).toBe('first');
    expect(chunks[1][FIRST_CHUNK_ERROR_KEY]).toBe(true);
  });

  it('should produce stop:abort SSE event through full pipeline when request is aborted', async () => {
    async function* abortingStream() {
      yield { type: 'message_start', message: { id: 'msg_1', content: [] } };
      throw new Error('Request was aborted.');
    }

    const identity = (chunk: any) => ({ data: chunk, id: 'msg_1', type: 'data' as const });
    const readable = convertIterableToStream(abortingStream())
      .pipeThrough(createTokenSpeedCalculator(identity))
      .pipeThrough(createSSEProtocolTransformer((c) => c));

    const reader = readable.getReader();
    const chunks: string[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value as string);
    }

    // Last 3 chunks should be the stop:abort SSE event
    const stopLines = chunks.slice(-3);
    expect(stopLines).toEqual(['id: \n', 'event: stop\n', `data: ${JSON.stringify('abort')}\n\n`]);
  });
});

describe('readableFromAsyncIterable', () => {
  it('should emit ABORT_CHUNK when abort error occurs during pull', async () => {
    async function* abortingStream() {
      yield 'first';
      throw new Error('Request was aborted.');
    }

    const readable = readableFromAsyncIterable(abortingStream());
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks).toEqual(['first', ABORT_CHUNK]);
  });

  it('should still surface non-abort errors as error chunks', async () => {
    async function* erroringStream() {
      yield 'first';
      throw new Error('rate limit');
    }

    const readable = readableFromAsyncIterable(erroringStream()).pipeThrough(
      createFirstErrorHandleTransformer(),
    );
    const reader = readable.getReader();
    const chunks: any[] = [];

    while (true) {
      const { done, value } = await reader.read();
      if (done) break;
      chunks.push(value);
    }

    expect(chunks[0]).toBe('first');
    expect(chunks[1][FIRST_CHUNK_ERROR_KEY]).toBe(true);
    expect(chunks[1].message).toBe('rate limit');
  });
});

describe('createSSEProtocolTransformer', () => {
  const processChunk = async (transformer: TransformStream, chunk: any) => {
    const results: any[] = [];
    const readable = new ReadableStream({
      start(controller) {
        controller.enqueue(chunk);
        controller.close();
      },
    });

    const writable = new WritableStream({
      write(chunk) {
        results.push(chunk);
      },
    });

    await readable.pipeThrough(transformer).pipeTo(writable);

    return results;
  };

  it('should convert chunk into SSE formatted lines without enforcing terminal (default)', async () => {
    const transformerFn = (chunk: any) => ({ type: 'text', id: chunk.id, data: chunk.data });
    const transformer = createSSEProtocolTransformer(transformerFn as any);

    const input = { id: '1', data: 'hello' };
    const results = await processChunk(transformer, input);

    // Should only output the text event, no injected error on flush (default not enforced)
    expect(results).toEqual([`id: 1\n`, `event: text\n`, `data: ${JSON.stringify('hello')}\n\n`]);
  });

  it('should not emit flush error if a terminal event was received (enforced)', async () => {
    const transformerFn = (chunk: any) => ({ type: 'stop', id: chunk.id, data: chunk.data });
    const transformer = createSSEProtocolTransformer(
      transformerFn as any,
      { id: 'stream_ok' },
      { requireTerminalEvent: true },
    );

    const input = { id: 'ok', data: 'bye' };
    const results = await processChunk(transformer, input);

    // Only the stop event lines should be present (no extra error event from flush)
    expect(results).toEqual([`id: ok\n`, `event: stop\n`, `data: ${JSON.stringify('bye')}\n\n`]);
  });

  it('should emit an error event on flush when no terminal event received (enforced)', async () => {
    const transformerFn = (chunk: any) => ({ type: 'text', id: chunk.id, data: chunk.data });
    const streamStack = { id: 'stream_missing_term' } as any;
    const transformer = createSSEProtocolTransformer(transformerFn as any, streamStack, {
      requireTerminalEvent: true,
    });

    const input = { id: '1', data: 'partial' };
    const results = await processChunk(transformer, input);

    // original 3 lines + 3 lines from flush error
    expect(results).toHaveLength(6);

    // last three lines should be the injected error event
    const lastThree = results.slice(-3);
    const expectedData = {
      body: { name: 'Stream parsing error', reason: 'unexpected_end' },
      message: 'Stream ended unexpectedly',
      name: 'Stream parsing error',
      type: 'StreamChunkError',
    };

    expect(lastThree).toEqual([
      `id: ${streamStack.id}\n`,
      `event: error\n`,
      `data: ${JSON.stringify(expectedData)}\n\n`,
    ]);
  });
});

describe('createCallbacksTransformer', () => {
  // Helper function to process chunks through transformer
  const processChunks = async (
    transformer: TransformStream<string, Uint8Array>,
    chunks: string[],
  ) => {
    const results: Uint8Array[] = [];
    const readable = new ReadableStream<string>({
      start(controller) {
        for (const chunk of chunks) {
          controller.enqueue(chunk);
        }
        controller.close();
      },
    });

    const writable = new WritableStream<Uint8Array>({
      write(chunk) {
        results.push(chunk);
      },
    });

    await readable.pipeThrough(transformer).pipeTo(writable);

    return results;
  };

  it('should call onStart callback when stream starts', async () => {
    const onStart = vi.fn();
    const transformer = createCallbacksTransformer({ onStart });

    await processChunks(transformer, []);

    expect(onStart).toHaveBeenCalledOnce();
  });

  it('should handle text chunks and call onText callback', async () => {
    const onText = vi.fn();
    const transformer = createCallbacksTransformer({ onText });

    const chunks = ['event: text\n', 'data: "Hello"\n\n', 'event: text\n', 'data: " World"\n\n'];

    await processChunks(transformer, chunks);

    expect(onText).toHaveBeenCalledTimes(2);
    expect(onText).toHaveBeenNthCalledWith(1, 'Hello');
    expect(onText).toHaveBeenNthCalledWith(2, ' World');
  });

  it('should handle reasoning chunks and call onThinking callback', async () => {
    const onThinking = vi.fn();
    const transformer = createCallbacksTransformer({ onThinking });

    const chunks = [
      'event: reasoning\n',
      'data: "Let me think..."\n\n',
      'event: reasoning\n',
      'data: " about this"\n\n',
    ];

    await processChunks(transformer, chunks);

    expect(onThinking).toHaveBeenCalledTimes(2);
    expect(onThinking).toHaveBeenNthCalledWith(1, 'Let me think...');
    expect(onThinking).toHaveBeenNthCalledWith(2, ' about this');
  });

  it('should handle base64_image chunks and call onBase64Image callback', async () => {
    const receivedCalls: Array<{
      image: { id: string; data: string };
      images: Array<{ id: string; data: string }>;
    }> = [];
    const onBase64Image = vi.fn((data) => {
      // Clone the data to capture the state at call time
      receivedCalls.push({
        image: { ...data.image },
        images: [...data.images],
      });
    });
    const transformer = createCallbacksTransformer({ onBase64Image });

    const imageData1 = { image: { id: 'img1', data: 'base64data1' }, images: [] };
    const imageData2 = { image: { id: 'img2', data: 'base64data2' }, images: [] };

    const chunks = [
      'event: base64_image\n',
      `data: ${JSON.stringify(imageData1)}\n\n`,
      'event: base64_image\n',
      `data: ${JSON.stringify(imageData2)}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onBase64Image).toHaveBeenCalledTimes(2);
    // Check the captured state at each call time
    expect(receivedCalls[0]).toEqual({
      image: { id: 'img1', data: 'base64data1' },
      images: [{ id: 'img1', data: 'base64data1' }],
    });
    expect(receivedCalls[1]).toEqual({
      image: { id: 'img2', data: 'base64data2' },
      images: [
        { id: 'img1', data: 'base64data1' },
        { id: 'img2', data: 'base64data2' },
      ],
    });
  });

  it('should handle content_part chunks and call onContentPart callback', async () => {
    const onContentPart = vi.fn();
    const transformer = createCallbacksTransformer({ onContentPart });

    const partData = {
      content: 'Hello',
      partType: 'text',
      mimeType: 'text/plain',
      thoughtSignature: 'sig123',
    };

    const chunks = ['event: content_part\n', `data: ${JSON.stringify(partData)}\n\n`];

    await processChunks(transformer, chunks);

    expect(onContentPart).toHaveBeenCalledOnce();
    expect(onContentPart).toHaveBeenCalledWith({
      content: 'Hello',
      partType: 'text',
      mimeType: 'text/plain',
      thoughtSignature: 'sig123',
    });
  });

  it('should handle reasoning_part chunks and call onReasoningPart callback', async () => {
    const onReasoningPart = vi.fn();
    const transformer = createCallbacksTransformer({ onReasoningPart });

    const partData = {
      content: 'base64ImageData',
      partType: 'image',
      mimeType: 'image/png',
      inReasoning: true,
    };

    const chunks = ['event: reasoning_part\n', `data: ${JSON.stringify(partData)}\n\n`];

    await processChunks(transformer, chunks);

    expect(onReasoningPart).toHaveBeenCalledOnce();
    expect(onReasoningPart).toHaveBeenCalledWith({
      content: 'base64ImageData',
      partType: 'image',
      mimeType: 'image/png',
    });
  });

  it('should handle usage chunks and call onUsage callback', async () => {
    const onUsage = vi.fn();
    const transformer = createCallbacksTransformer({ onUsage });

    const usageData = {
      inputTextTokens: 10,
      outputTextTokens: 20,
      totalTokens: 30,
    };

    const chunks = ['event: usage\n', `data: ${JSON.stringify(usageData)}\n\n`];

    await processChunks(transformer, chunks);

    expect(onUsage).toHaveBeenCalledOnce();
    expect(onUsage).toHaveBeenCalledWith(usageData);
  });

  it('should handle grounding chunks and call onGrounding callback', async () => {
    const onGrounding = vi.fn();
    const transformer = createCallbacksTransformer({ onGrounding });

    const groundingData = {
      sources: [{ url: 'https://example.com', title: 'Example' }],
    };

    const chunks = ['event: grounding\n', `data: ${JSON.stringify(groundingData)}\n\n`];

    await processChunks(transformer, chunks);

    expect(onGrounding).toHaveBeenCalledOnce();
    expect(onGrounding).toHaveBeenCalledWith(groundingData);
  });

  it('should handle tool_calls chunks and call onToolsCalling callback', async () => {
    const onToolsCalling = vi.fn();
    const transformer = createCallbacksTransformer({ onToolsCalling });

    const toolCallData = [
      {
        index: 0,
        id: 'call_123',
        type: 'function',
        function: { name: 'get_weather', arguments: '{"city":"SF"}' },
      },
    ];

    const chunks = ['event: tool_calls\n', `data: ${JSON.stringify(toolCallData)}\n\n`];

    await processChunks(transformer, chunks);

    expect(onToolsCalling).toHaveBeenCalledOnce();
    expect(onToolsCalling).toHaveBeenCalledWith({
      chunk: toolCallData,
      toolsCalling: expect.any(Array),
    });
  });

  it('should call onCompletion and onFinal callbacks on flush with aggregated data', async () => {
    const onCompletion = vi.fn();
    const onFinal = vi.fn();
    const transformer = createCallbacksTransformer({ onCompletion, onFinal });

    const chunks = [
      'event: text\n',
      'data: "Hello"\n\n',
      'event: text\n',
      'data: " World"\n\n',
      'event: reasoning\n',
      'data: "Thinking..."\n\n',
      'event: usage\n',
      `data: ${JSON.stringify({ totalTokens: 10 })}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onCompletion).toHaveBeenCalledOnce();
    expect(onFinal).toHaveBeenCalledOnce();

    const expectedData = {
      text: 'Hello World',
      thinking: 'Thinking...',
      usage: { totalTokens: 10 },
      grounding: undefined,
      speed: undefined,
      toolsCalling: undefined,
    };

    expect(onCompletion).toHaveBeenCalledWith(expectedData);
    expect(onFinal).toHaveBeenCalledWith(expectedData);
  });

  it('should capture finishReason from stop chunks and include in final data', async () => {
    const onCompletion = vi.fn();
    const onFinal = vi.fn();
    const transformer = createCallbacksTransformer({ onCompletion, onFinal });

    // Simulates the Gemini "soft interrupt" path: empty content + non-STOP finishReason
    // (e.g. RECITATION / MAX_TOKENS) — we MUST capture the reason so downstream
    // tracing/UI can surface it instead of silently rendering empty.
    const chunks = [
      'event: stop\n',
      `data: ${JSON.stringify('RECITATION')}\n\n`,
      'event: usage\n',
      `data: ${JSON.stringify({ totalTokens: 10 })}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onCompletion).toHaveBeenCalledWith(
      expect.objectContaining({ finishReason: 'RECITATION' }),
    );
    expect(onFinal).toHaveBeenCalledWith(expect.objectContaining({ finishReason: 'RECITATION' }));
  });

  it('should keep the first finishReason when multiple stop chunks are emitted', async () => {
    // Anthropic emits message_delta (carrying real stop_reason) followed by a
    // message_stop sentinel — the meaningful reason must not be clobbered.
    const onCompletion = vi.fn();
    const transformer = createCallbacksTransformer({ onCompletion });

    const chunks = [
      'event: stop\n',
      `data: ${JSON.stringify('max_tokens')}\n\n`,
      'event: stop\n',
      `data: ${JSON.stringify('message_stop')}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onCompletion).toHaveBeenCalledWith(
      expect.objectContaining({ finishReason: 'max_tokens' }),
    );
  });

  it('should fall back to a later stop chunk when the first one is empty', async () => {
    const onCompletion = vi.fn();
    const transformer = createCallbacksTransformer({ onCompletion });

    const chunks = [
      'event: stop\n',
      `data: ${JSON.stringify('')}\n\n`,
      'event: stop\n',
      `data: ${JSON.stringify('end_turn')}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onCompletion).toHaveBeenCalledWith(
      expect.objectContaining({ finishReason: 'end_turn' }),
    );
  });

  it('should leave finishReason undefined when no stop chunk is received', async () => {
    const onFinal = vi.fn();
    const transformer = createCallbacksTransformer({ onFinal });

    const chunks = ['event: text\n', 'data: "Hi"\n\n'];

    await processChunks(transformer, chunks);

    expect(onFinal).toHaveBeenCalledWith(expect.objectContaining({ finishReason: undefined }));
  });

  it('should handle speed chunks and include in final data', async () => {
    const onFinal = vi.fn();
    const transformer = createCallbacksTransformer({ onFinal });

    const speedData = { tps: 50, ttft: 100, duration: 500, latency: 600 };

    const chunks = [
      'event: text\n',
      'data: "Hi"\n\n',
      'event: speed\n',
      `data: ${JSON.stringify(speedData)}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onFinal).toHaveBeenCalledWith(
      expect.objectContaining({
        speed: speedData,
      }),
    );
  });

  it('should work without any callbacks', async () => {
    const transformer = createCallbacksTransformer(undefined);

    const chunks = ['event: text\n', 'data: "Hello"\n\n'];

    // Should not throw
    await expect(processChunks(transformer, chunks)).resolves.toBeDefined();
  });

  it('should ignore invalid JSON data gracefully', async () => {
    const onText = vi.fn();
    const transformer = createCallbacksTransformer({ onText });

    const chunks = [
      'event: text\n',
      'data: invalid-json\n\n',
      'event: text\n',
      'data: "Valid"\n\n',
    ];

    await processChunks(transformer, chunks);

    // Only the valid JSON should trigger callback
    expect(onText).toHaveBeenCalledOnce();
    expect(onText).toHaveBeenCalledWith('Valid');
  });

  it('should handle multiple tool_calls chunks and accumulate them', async () => {
    const onToolsCalling = vi.fn();
    const transformer = createCallbacksTransformer({ onToolsCalling });

    const toolCall1 = [
      {
        index: 0,
        id: 'call_1',
        type: 'function',
        function: { name: 'func1', arguments: '' },
      },
    ];

    const toolCall2 = [
      {
        index: 0,
        function: { arguments: '{"a":1}' },
      },
    ];

    const chunks = [
      'event: tool_calls\n',
      `data: ${JSON.stringify(toolCall1)}\n\n`,
      'event: tool_calls\n',
      `data: ${JSON.stringify(toolCall2)}\n\n`,
    ];

    await processChunks(transformer, chunks);

    expect(onToolsCalling).toHaveBeenCalledTimes(2);
  });

  // Regression: stream errors silently swallowed by createCallbacksTransformer
  // These tests assert the CORRECT expected behavior. They will FAIL until the bug is fixed.
  describe('error event handling', () => {
    it('should call onError callback when stream contains an error event', async () => {
      const onError = vi.fn();
      const onText = vi.fn();
      const onCompletion = vi.fn();
      const transformer = createCallbacksTransformer({ onCompletion, onError, onText } as any);

      const errorPayload = {
        body: { message: 'rate limit exceeded' },
        message: 'rate limit exceeded',
        type: 'ProviderBizError',
      };

      const chunks = ['event: error\n', `data: ${JSON.stringify(errorPayload)}\n\n`];

      await processChunks(transformer, chunks);

      // onText should NOT be called
      expect(onText).not.toHaveBeenCalled();

      // onError SHOULD be called with the error data
      expect(onError).toHaveBeenCalledOnce();
      expect(onError).toHaveBeenCalledWith(errorPayload);
    });

    it('should include error in onCompletion data when stream has error after partial text', async () => {
      const onCompletion = vi.fn();
      const transformer = createCallbacksTransformer({ onCompletion } as any);

      const errorPayload = {
        body: { message: 'content filter triggered' },
        message: 'content filter triggered',
        type: 'ProviderBizError',
      };

      const chunks = [
        'event: text\n',
        'data: "Partial response"\n\n',
        'event: error\n',
        `data: ${JSON.stringify(errorPayload)}\n\n`,
      ];

      await processChunks(transformer, chunks);

      // onCompletion should include the error so callers can detect the failure
      expect(onCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          error: errorPayload,
          text: 'Partial response',
        }),
      );
    });

    it('should surface first-chunk error via onError callback', async () => {
      // Simulates the full chain: provider throws → ERROR_CHUNK_PREFIX → FIRST_CHUNK_ERROR_KEY
      // → transformOpenAIStream returns { type: 'error' } → createSSEProtocolTransformer
      // → createCallbacksTransformer should handle 'error' in switch
      const onError = vi.fn();
      const onCompletion = vi.fn();
      const transformer = createCallbacksTransformer({ onCompletion, onError } as any);

      const errorPayload = {
        body: { message: 'insufficient balance', status_code: 1008 },
        message: 'insufficient balance',
        type: 'ProviderBizError',
      };

      const chunks = ['event: error\n', `data: ${JSON.stringify(errorPayload)}\n\n`];

      await processChunks(transformer, chunks);

      // onError should be called
      expect(onError).toHaveBeenCalledOnce();
      expect(onError).toHaveBeenCalledWith(errorPayload);

      // onCompletion should include the error information
      expect(onCompletion).toHaveBeenCalledWith(
        expect.objectContaining({
          error: errorPayload,
        }),
      );
    });
  });
});
