import { describe, test, expect } from 'bun:test';
import { processLine, createParseState, parseStream } from './stream-parser';
import type { ChatEvent } from '../../api/chat/types';
import type { StreamParserCallbacks, ParseState } from './stream-parser';

/** The main agent's accumulated deltas — buffers are keyed by parent_tool_use_id, '' being the agent. */
const buffered = (state: ParseState, parent = '') => state.textBuffers.get(parent);

function makeCallbacks(): { events: ChatEvent[]; sessionIds: string[]; callbacks: StreamParserCallbacks } {
  const events: ChatEvent[] = [];
  const sessionIds: string[] = [];
  return {
    events,
    sessionIds,
    callbacks: {
      onEvent: (e) => events.push(e),
      onSessionId: (id) => sessionIds.push(id),
    },
  };
}

describe('processLine', () => {
  test('skips empty lines', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    expect(processLine('', state, callbacks)).toBe(false);
    expect(processLine('   ', state, callbacks)).toBe(false);
    expect(events).toHaveLength(0);
  });

  test('skips malformed JSON', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    expect(processLine('not json', state, callbacks)).toBe(false);
    expect(events).toHaveLength(0);
  });

  test('handles stream_event text delta', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'stream_event',
      event: { type: 'content_block_delta', delta: { type: 'text_delta', text: 'hello' } },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([{ type: 'delta', text: 'hello' }]);
    expect(buffered(state)).toBe('hello');
  });

  test('accumulates text buffer across deltas', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const mkDelta = (text: string) =>
      JSON.stringify({
        type: 'stream_event',
        event: { type: 'content_block_delta', delta: { type: 'text_delta', text } },
      });
    processLine(mkDelta('hello '), state, callbacks);
    processLine(mkDelta('world'), state, callbacks);
    expect(buffered(state)).toBe('hello world');
    expect(events).toHaveLength(2);
  });

  test('handles assistant text block — clears text buffer', () => {
    const state = createParseState();
    state.textBuffers.set('', 'partial');
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'assistant',
      message: { content: [{ type: 'text', text: 'full response' }] },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([{ type: 'text', text: 'full response' }]);
    expect(buffered(state)).toBeUndefined();
  });

  test('handles assistant tool_use block — flushes text buffer first', () => {
    const state = createParseState();
    state.textBuffers.set('', 'thinking...');
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'assistant',
      message: {
        content: [{ type: 'tool_use', id: 'call_1', name: 'sqlite', input: { query: 'SELECT 1' } }],
      },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([
      { type: 'text', text: 'thinking...' },
      { type: 'tool:start', toolCallId: 'call_1', toolName: 'sqlite', toolInput: { query: 'SELECT 1' } },
    ]);
    expect(buffered(state)).toBeUndefined();
  });

  test('handles assistant tool_use without prior text buffer', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'assistant',
      message: {
        content: [{ type: 'tool_use', id: 'call_1', name: 'email_db', input: {} }],
      },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([{ type: 'tool:start', toolCallId: 'call_1', toolName: 'email_db', toolInput: {} }]);
  });

  test('handles user tool_result with string content', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'user',
      message: {
        content: [{ type: 'tool_result', tool_use_id: 'call_1', content: 'result text', is_error: false }],
      },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([{ type: 'tool:result', toolCallId: 'call_1', output: 'result text', isError: false }]);
  });

  test('handles user tool_result with array content', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'user',
      message: {
        content: [
          {
            type: 'tool_result',
            tool_use_id: 'call_2',
            content: [
              { type: 'text', text: 'line 1' },
              { type: 'image', data: 'ignored' },
              { type: 'text', text: 'line 2' },
            ],
            is_error: true,
          },
        ],
      },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([{ type: 'tool:result', toolCallId: 'call_2', output: 'line 1\nline 2', isError: true }]);
  });

  test('handles system init — captures session id', () => {
    const state = createParseState();
    const { sessionIds, callbacks } = makeCallbacks();
    const line = JSON.stringify({ type: 'system', subtype: 'init', session_id: 'sess_abc123' });
    processLine(line, state, callbacks);
    expect(sessionIds).toEqual(['sess_abc123']);
  });

  test('handles compact_boundary — emits compact:done with what it dropped', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'system',
      subtype: 'compact_boundary',
      compact_metadata: { trigger: 'auto', pre_tokens: 168215 },
    });
    processLine(line, state, callbacks);
    expect(events).toEqual([{ type: 'compact:done', trigger: 'auto', preTokens: 168215 }]);
  });

  test('handles compact_boundary with no metadata — still announces the boundary', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    processLine(JSON.stringify({ type: 'system', subtype: 'compact_boundary' }), state, callbacks);
    expect(events).toEqual([{ type: 'compact:done', trigger: 'auto', preTokens: 0 }]);
  });

  test('handles result — sets gotResult, emits cost', () => {
    const state = createParseState();
    state.textBuffers.set('', 'trailing');
    const { events, sessionIds, callbacks } = makeCallbacks();
    const line = JSON.stringify({
      type: 'result',
      is_error: false,
      result: 'done',
      session_id: 'sess_xyz',
      usage: { input_tokens: 100, output_tokens: 50 },
      total_cost_usd: 0.003,
    });
    processLine(line, state, callbacks);
    expect(state.gotResult).toBe(true);
    expect(sessionIds).toEqual(['sess_xyz']);
    expect(events).toEqual([
      { type: 'text', text: 'trailing' },
      { type: 'result', cost: { inputTokens: 100, outputTokens: 50, totalUSD: 0.003 }, claudeSessionId: 'sess_xyz' },
    ]);
  });

  test('handles error result — flushes buffer and emits error', () => {
    const state = createParseState();
    state.textBuffers.set('', 'partial');
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({ type: 'result', is_error: true, result: 'something broke' });
    processLine(line, state, callbacks);
    expect(state.gotResult).toBe(true);
    expect(events).toEqual([
      { type: 'text', text: 'partial' },
      { type: 'error', message: 'something broke' },
    ]);
  });

  test('handles result with missing usage — defaults to zero', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const line = JSON.stringify({ type: 'result', is_error: false });
    processLine(line, state, callbacks);
    expect(events).toEqual([
      { type: 'result', cost: { inputTokens: 0, outputTokens: 0, totalUSD: 0 }, claudeSessionId: undefined },
    ]);
  });

  // ── Subagent attribution ──
  //
  // The harness stamps every message a subagent produces with `parent_tool_use_id` — the id of the Task
  // tool call that spawned it. Nothing read it, so a subagent's prose arrived as a top-level assistant
  // message and its tool calls looked like the main agent's. Worse, its deltas landed in the one shared
  // text buffer, so a flush emitted both voices concatenated into a single message.

  test('stamps subagent output with its parent Task call', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    processLine(
      JSON.stringify({
        type: 'assistant',
        parent_tool_use_id: 'call_task',
        message: { content: [{ type: 'text', text: 'searched 40 files' }] },
      }),
      state,
      callbacks,
    );
    processLine(
      JSON.stringify({
        type: 'user',
        parent_tool_use_id: 'call_task',
        message: { content: [{ type: 'tool_result', tool_use_id: 'call_grep', content: 'hit', is_error: false }] },
      }),
      state,
      callbacks,
    );
    expect(events).toEqual([
      { type: 'text', text: 'searched 40 files', parentToolUseId: 'call_task' },
      { type: 'tool:result', toolCallId: 'call_grep', output: 'hit', isError: false, parentToolUseId: 'call_task' },
    ]);
  });

  test('the main agent and a subagent buffer their deltas separately', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    const delta = (text: string, parent?: string) =>
      JSON.stringify({
        type: 'stream_event',
        ...(parent ? { parent_tool_use_id: parent } : {}),
        event: { type: 'content_block_delta', delta: { type: 'text_delta', text } },
      });

    processLine(delta('I will '), state, callbacks);
    processLine(delta('reading the ', 'call_task'), state, callbacks);
    processLine(delta('delegate.'), state, callbacks);
    processLine(delta('config now.', 'call_task'), state, callbacks);

    expect(buffered(state)).toBe('I will delegate.');
    expect(buffered(state, 'call_task')).toBe('reading the config now.');

    // And a tool call by the subagent flushes only the subagent's sentence.
    processLine(
      JSON.stringify({
        type: 'assistant',
        parent_tool_use_id: 'call_task',
        message: { content: [{ type: 'tool_use', id: 'call_read', name: 'Read', input: {} }] },
      }),
      state,
      callbacks,
    );
    expect(events.at(-2)).toEqual({ type: 'text', text: 'reading the config now.', parentToolUseId: 'call_task' });
    expect(buffered(state)).toBe('I will delegate.');
  });

  test('a result flushes every buffer, subagents included', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    state.textBuffers.set('', 'main said this');
    state.textBuffers.set('call_task', 'subagent said this');

    processLine(JSON.stringify({ type: 'result', is_error: false }), state, callbacks);

    expect(events.slice(0, 2)).toEqual([
      { type: 'text', text: 'main said this' },
      { type: 'text', text: 'subagent said this', parentToolUseId: 'call_task' },
    ]);
  });

  test('ignores unknown message types', () => {
    const state = createParseState();
    const { events, callbacks } = makeCallbacks();
    expect(processLine(JSON.stringify({ type: 'unknown_thing' }), state, callbacks)).toBe(false);
    expect(events).toHaveLength(0);
  });
});

describe('parseStream', () => {
  function makeStream(lines: string[]): ReadableStream<Uint8Array> {
    const encoder = new TextEncoder();
    const data = encoder.encode(lines.join('\n') + '\n');
    return new ReadableStream({
      start(controller) {
        controller.enqueue(data);
        controller.close();
      },
    });
  }

  test('parses a full conversation stream', async () => {
    const lines = [
      JSON.stringify({ type: 'system', subtype: 'init', session_id: 'sess_1' }),
      JSON.stringify({
        type: 'stream_event',
        event: { type: 'content_block_delta', delta: { type: 'text_delta', text: 'Hello' } },
      }),
      JSON.stringify({
        type: 'assistant',
        message: { content: [{ type: 'text', text: 'Hello there' }] },
      }),
      JSON.stringify({
        type: 'result',
        is_error: false,
        session_id: 'sess_1',
        usage: { input_tokens: 10, output_tokens: 5 },
        total_cost_usd: 0.001,
      }),
    ];

    const { events, sessionIds, callbacks } = makeCallbacks();
    const state = await parseStream(makeStream(lines), callbacks);

    expect(state.gotResult).toBe(true);
    expect(sessionIds).toEqual(['sess_1', 'sess_1']);
    // `session` leads: the transcript id goes out at `system.init` so the URL is a permalink from the
    // start of the turn, which is what makes a mid-turn refresh reattachable.
    expect(events.map((e) => e.type)).toEqual(['session', 'delta', 'text', 'result']);
    expect(events[0]).toEqual({ type: 'session', claudeSessionId: 'sess_1' });
  });

  test('handles chunked delivery (split mid-line)', async () => {
    const fullLine = JSON.stringify({
      type: 'stream_event',
      event: { type: 'content_block_delta', delta: { type: 'text_delta', text: 'chunked' } },
    });
    const resultLine = JSON.stringify({ type: 'result', is_error: false });

    const encoder = new TextEncoder();
    const mid = Math.floor(fullLine.length / 2);
    const chunk1 = encoder.encode(fullLine.slice(0, mid));
    const chunk2 = encoder.encode(fullLine.slice(mid) + '\n' + resultLine + '\n');

    const stream = new ReadableStream({
      start(controller) {
        controller.enqueue(chunk1);
        controller.enqueue(chunk2);
        controller.close();
      },
    });

    const { events, callbacks } = makeCallbacks();
    const state = await parseStream(stream, callbacks);

    expect(state.gotResult).toBe(true);
    expect(events[0]).toEqual({ type: 'delta', text: 'chunked' });
  });

  test('handles trailing data without newline', async () => {
    const line = JSON.stringify({ type: 'result', is_error: false });
    const encoder = new TextEncoder();
    const stream = new ReadableStream({
      start(controller) {
        controller.enqueue(encoder.encode(line)); // no trailing newline
        controller.close();
      },
    });

    const { events, callbacks } = makeCallbacks();
    const state = await parseStream(stream, callbacks);
    expect(state.gotResult).toBe(true);
  });
});