import type {
  OllamaChatOptions,
  OllamaGenerateOptions,
  OllamaStreamChunk,
  OllamaGenerateChunk,
} from './types.js';
import { parseNdjsonStream } from './ndjson.js';

/**
 * Stream a chat completion from Ollama.
 *
 * Yields NDJSON chunks from `POST /api/chat` as an async generator.
 *
 * @param baseUrl - Ollama server base URL
 * @param options - Chat options (model, messages, signal, etc.)
 */
export async function* streamChat(
  baseUrl: string,
  options: OllamaChatOptions
): AsyncGenerator<OllamaStreamChunk> {
  const { model, messages, signal, ...rest } = options;

  const body: Record<string, unknown> = { model, messages, stream: true };
  if (rest.options) body.options = rest.options;
  if (rest.format) body.format = rest.format;
  if (rest.keep_alive) body.keep_alive = rest.keep_alive;

  const res = await fetch(`${baseUrl}/api/chat`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(body),
    signal,
  });

  if (!res.ok) {
    const text = await res.text().catch(() => '');
    throw new Error(`Ollama chat failed (${res.status}): ${text.slice(0, 200)}`);
  }

  if (!res.body) {
    throw new Error('No response body from Ollama');
  }

  yield* parseNdjsonStream<OllamaStreamChunk>(res.body);
}

/**
 * Stream a text generation from Ollama.
 *
 * Yields NDJSON chunks from `POST /api/generate` as an async generator.
 *
 * @param baseUrl - Ollama server base URL
 * @param options - Generate options (model, prompt, signal, etc.)
 */
export async function* streamGenerate(
  baseUrl: string,
  options: OllamaGenerateOptions
): AsyncGenerator<OllamaGenerateChunk> {
  const { model, prompt, signal, ...rest } = options;

  const body: Record<string, unknown> = { model, prompt, stream: true };
  if (rest.system) body.system = rest.system;
  if (rest.options) body.options = rest.options;
  if (rest.format) body.format = rest.format;
  if (rest.keep_alive) body.keep_alive = rest.keep_alive;
  if (rest.context) body.context = rest.context;

  const res = await fetch(`${baseUrl}/api/generate`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(body),
    signal,
  });

  if (!res.ok) {
    const text = await res.text().catch(() => '');
    throw new Error(`Ollama generate failed (${res.status}): ${text.slice(0, 200)}`);
  }

  if (!res.body) {
    throw new Error('No response body from Ollama');
  }

  yield* parseNdjsonStream<OllamaGenerateChunk>(res.body);
}