import { describe, it, expect } from 'vitest';
import { runIsolatedLlm, consumeLlmStream } from './llm-stream.js';
import type { LLMEvent, Message, ToolDef } from '../llm/openai-compat.js';

class FakeClient {
  readonly calls: Array<{ messages: unknown; tools?: unknown }> = [];
  private index = 0;
  readonly timeoutMs = 60_000;

  constructor(private readonly responses: LLMEvent[][]) {}

  async *chat(messages: unknown, tools?: unknown, _signal?: AbortSignal): AsyncGenerator<LLMEvent> {
    this.calls.push({ messages, tools });
    const response = this.responses[this.index++] ?? [];
    for (const event of response) yield event;
  }
}

const NO_TOOLS: ToolDef[] = [];
const SHORT_TIMEOUT = 5_000;

describe('runIsolatedLlm', () => {
  it('concatenates text events into the returned string', async () => {
    const client = new FakeClient([[
      { type: 'text', text: 'hello ' },
      { type: 'text', text: 'world' },
      { type: 'done' },
    ]]);
    const out = await runIsolatedLlm(client as never, [{ role: 'user', content: 'hi' }]);
    expect(out).toBe('hello world');
  });

  it('strips thinking tokens from the result', async () => {
    const client = new FakeClient([[
      { type: 'text', text: '<think>private reasoning</think>final answer' },
      { type: 'done' },
    ]]);
    const out = await runIsolatedLlm(client as never, [{ role: 'user', content: 'hi' }]);
    expect(out).toBe('final answer');
  });

  it('throws when the LLM tries to invoke a tool', async () => {
    const client = new FakeClient([[
      { type: 'tool_use', id: 'x', name: 'Read', input: { file_path: '/x' } },
      { type: 'done' },
    ]]);
    await expect(runIsolatedLlm(client as never, [{ role: 'user', content: 'hi' }]))
      .rejects.toThrow(/unexpectedly requested tool "Read"/);
  });

  it('throws when an error event is yielded', async () => {
    const client = new FakeClient([[
      { type: 'error', error: 'rate limited' },
    ]]);
    await expect(runIsolatedLlm(client as never, [{ role: 'user', content: 'hi' }]))
      .rejects.toThrow('rate limited');
  });
});

describe('consumeLlmStream', () => {
  it('accumulates text and tool calls separately', async () => {
    const client = new FakeClient([[
      { type: 'text', text: 'thinking out loud, ' },
      { type: 'tool_use', id: 't1', name: 'Read', input: { file_path: '/a' } },
      { type: 'text', text: 'and more text' },
      { type: 'tool_use', id: 't2', name: 'Glob', input: { pattern: '*.ts' } },
      { type: 'done', usage: { prompt_tokens: 100, completion_tokens: 20 } },
    ]]);
    const result = await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT);
    expect(result.accumulatedText).toBe('thinking out loud, and more text');
    expect(result.pendingToolCalls).toHaveLength(2);
    expect(result.pendingToolCalls[0]!.function.name).toBe('Read');
    expect(result.pendingToolCalls[1]!.function.name).toBe('Glob');
    expect(result.pendingToolCalls[0]!.function.arguments).toBe('{"file_path":"/a"}');
    expect(result.lastUsage).toEqual({ prompt_tokens: 100, completion_tokens: 20 });
    expect(result.hadError).toBe(false);
  });

  it('strips thinking tokens from the accumulated text', async () => {
    const client = new FakeClient([[
      { type: 'text', text: '<think>internal</think>visible output' },
      { type: 'done' },
    ]]);
    const result = await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT);
    expect(result.accumulatedText).toBe('visible output');
  });

  it('captures error events without throwing', async () => {
    const client = new FakeClient([[
      { type: 'error', error: 'context length exceeded' },
    ]]);
    const result = await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT);
    expect(result.hadError).toBe(true);
    expect(result.errorMessage).toBe('context length exceeded');
  });

  it('fires onText / onToolUse callbacks for each event', async () => {
    const textChunks: string[] = [];
    const toolUses: Array<{ name: string; input: Record<string, unknown> }> = [];
    const client = new FakeClient([[
      { type: 'text', text: 'a' },
      { type: 'text', text: 'b' },
      { type: 'tool_use', id: 't1', name: 'Read', input: { file_path: '/x' } },
      { type: 'done' },
    ]]);
    await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT, {
      onText: (text) => textChunks.push(text),
      onToolUse: (name, input) => toolUses.push({ name, input }),
    });
    expect(textChunks).toEqual(['a', 'b']);
    expect(toolUses).toEqual([{ name: 'Read', input: { file_path: '/x' } }]);
  });

  it('fires onToolCallDelta for each tool_use_delta without adding to pendingToolCalls', async () => {
    const deltas: Array<{ index: number; callId: string; name: string; chunk: string }> = [];
    const client = new FakeClient([[
      { type: 'tool_use_delta', index: 0, callId: 'c1', name: 'Write', chunk: '{"content":"a' },
      { type: 'tool_use_delta', index: 0, callId: 'c1', name: 'Write', chunk: 'b"}' },
      { type: 'tool_use', id: 'c1', name: 'Write', input: { content: 'ab' } },
      { type: 'done' },
    ]]);
    const result = await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT, {
      onToolCallDelta: (index, callId, name, chunk) => deltas.push({ index, callId, name, chunk }),
    });
    expect(deltas).toEqual([
      { index: 0, callId: 'c1', name: 'Write', chunk: '{"content":"a' },
      { index: 0, callId: 'c1', name: 'Write', chunk: 'b"}' },
    ]);
    expect(result.pendingToolCalls).toHaveLength(1); // only the aggregated tool_use
    expect(result.pendingToolCalls[0]!.function.name).toBe('Write');
  });

  it('captures backend events from proxy clients and surfaces via callback + result', async () => {
    const backends: Array<{ id: string; cacheKey: string | null }> = [];
    const client = new FakeClient([[
      { type: 'backend', backendId: 'gpu-rtx-a', cacheKey: null },
      { type: 'text', text: 'ok' },
      { type: 'done' },
    ]]);
    const result = await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT, {
      onBackend: (backendId, cacheKey) => backends.push({ id: backendId, cacheKey }),
    });
    expect(backends).toEqual([{ id: 'gpu-rtx-a', cacheKey: null }]);
    expect(result.backendId).toBe('gpu-rtx-a');
    expect(result.backendCacheKey).toBeNull();
    expect(result.accumulatedText).toBe('ok');
  });

  it('records backend cacheKey when present (LiteLLM cache hit)', async () => {
    const client = new FakeClient([[
      { type: 'backend', backendId: 'gpu-h100-b', cacheKey: 'sha:cached' },
      { type: 'text', text: 'cached' },
      { type: 'done' },
    ]]);
    const result = await consumeLlmStream(client as never, [], NO_TOOLS, undefined, SHORT_TIMEOUT);
    expect(result.backendId).toBe('gpu-h100-b');
    expect(result.backendCacheKey).toBe('sha:cached');
  });

  it('hits the idle safety timeout when the stream stalls', async () => {
    class StallingClient {
      readonly timeoutMs = 60_000;
      // eslint-disable-next-line @typescript-eslint/no-unused-vars, require-yield
      async *chat(_m: Message[], _t?: ToolDef[], _signal?: AbortSignal): AsyncGenerator<LLMEvent> {
        await new Promise((resolve) => setTimeout(resolve, 1_000));
      }
    }
    const result = await consumeLlmStream(
      new StallingClient() as never,
      [],
      NO_TOOLS,
      undefined,
      50, // 50ms idle budget — much shorter than the stalling 1s
    );
    expect(result.hadError).toBe(true);
    expect(result.errorMessage).toContain('idle safety timeout');
  });
});
