import { describe, it, expect, vi, beforeEach } from 'vitest';
import { streamChat, streamGenerate } from './stream.js';

const BASE_URL = 'http://localhost:11434';

function createNdjsonResponse(chunks: object[]): Response {
  const encoder = new TextEncoder();
  const ndjson = chunks.map(c => JSON.stringify(c)).join('\n') + '\n';
  const stream = new ReadableStream<Uint8Array>({
    start(controller) {
      controller.enqueue(encoder.encode(ndjson));
      controller.close();
    },
  });
  return { ok: true, status: 200, body: stream } as unknown as Response;
}

describe('streamChat', () => {
  beforeEach(() => {
    vi.restoreAllMocks();
  });

  it('yields chat stream chunks', async () => {
    const chunks = [
      { model: 'llama3', message: { role: 'assistant', content: 'Hello' }, done: false },
      { model: 'llama3', message: { role: 'assistant', content: ' world' }, done: false },
      { model: 'llama3', message: { role: 'assistant', content: '' }, done: true, eval_count: 10 },
    ];
    globalThis.fetch = vi.fn().mockResolvedValue(createNdjsonResponse(chunks));

    const results = [];
    for await (const chunk of streamChat(BASE_URL, {
      model: 'llama3',
      messages: [{ role: 'user', content: 'Hi' }],
    })) {
      results.push(chunk);
    }

    expect(results).toHaveLength(3);
    expect(results[0].message.content).toBe('Hello');
    expect(results[2].done).toBe(true);
    expect(results[2].eval_count).toBe(10);
  });

  it('throws on non-ok response', async () => {
    globalThis.fetch = vi.fn().mockResolvedValue({
      ok: false,
      status: 500,
      text: () => Promise.resolve('internal error'),
    });

    const gen = streamChat(BASE_URL, {
      model: 'llama3',
      messages: [{ role: 'user', content: 'Hi' }],
    });
    await expect(gen.next()).rejects.toThrow('Ollama chat failed (500)');
  });

  it('throws when response has no body', async () => {
    globalThis.fetch = vi.fn().mockResolvedValue({
      ok: true,
      status: 200,
      body: null,
    });

    const gen = streamChat(BASE_URL, {
      model: 'llama3',
      messages: [{ role: 'user', content: 'Hi' }],
    });
    await expect(gen.next()).rejects.toThrow('No response body');
  });
});

describe('streamGenerate', () => {
  beforeEach(() => {
    vi.restoreAllMocks();
  });

  it('yields generate stream chunks', async () => {
    const chunks = [
      { model: 'llama3', response: 'Hello', done: false },
      { model: 'llama3', response: ' world', done: false },
      { model: 'llama3', response: '', done: true, eval_count: 5 },
    ];
    globalThis.fetch = vi.fn().mockResolvedValue(createNdjsonResponse(chunks));

    const results = [];
    for await (const chunk of streamGenerate(BASE_URL, {
      model: 'llama3',
      prompt: 'Say hello',
    })) {
      results.push(chunk);
    }

    expect(results).toHaveLength(3);
    expect(results[0].response).toBe('Hello');
    expect(results[2].done).toBe(true);
  });
});