// Author: Simon-Pierre Boucher
// Contact: contact@spboucher.ai
// Project: chat.spboucher.ai

// Parses the OpenRouter SSE wire stream into normalized ChatStreamEvents.

import { orChatStream } from "./client";
import { parseStreamChunk } from "./schemas";
import { AIError, normalizeHttpError, normalizeNetworkError } from "./errors";
import type { ChatStreamEvent, GenerationRequest, ModelDefinition, ModelGateway } from "./types";
import { fetchModels } from "./models";

/**
 * Stream a generation as normalized events. Cancellation propagates upstream
 * through request.signal — aborting truly stops the OpenRouter generation.
 */
export async function* streamGeneration(
  request: GenerationRequest,
  generationId: string
): AsyncGenerator<ChatStreamEvent> {
  const signal = request.signal ?? new AbortController().signal;

  const body: Record<string, unknown> = {
    model: request.model,
    messages: request.messages,
  };
  if (request.temperature !== undefined) body.temperature = request.temperature;
  if (request.maxTokens !== undefined) body.max_tokens = request.maxTokens;
  if (request.routing) Object.assign(body, request.routing);

  let res: Response;
  try {
    res = await orChatStream(body, signal);
  } catch (e) {
    const err = e instanceof AIError ? e : normalizeNetworkError(e, request.model);
    yield { type: "generation.error", message: err.message, retryable: err.retryable };
    return;
  }

  yield { type: "generation.start", generationId, model: request.model };

  const reader = res.body!.getReader();
  const decoder = new TextDecoder();
  let buffer = "";
  const openToolCalls = new Set<string>();

  try {
    for (;;) {
      const { done, value } = await reader.read();
      if (done) break;
      buffer += decoder.decode(value, { stream: true });

      // SSE frames are separated by newlines; each data line is "data: {json}" or "data: [DONE]"
      let newlineIdx: number;
      while ((newlineIdx = buffer.indexOf("\n")) !== -1) {
        const line = buffer.slice(0, newlineIdx).trim();
        buffer = buffer.slice(newlineIdx + 1);
        if (!line.startsWith("data:")) continue; // comments / keep-alives
        const payload = line.slice(5).trim();
        if (payload === "[DONE]") continue;

        const chunk = parseStreamChunk(payload);
        if (!chunk) continue;

        if (chunk.error) {
          const status = typeof chunk.error.code === "number" ? chunk.error.code : 500;
          const err = normalizeHttpError(status, chunk.error.message ?? "", request.model);
          yield { type: "generation.error", message: err.message, retryable: err.retryable };
          return;
        }

        const delta = chunk.choices?.[0]?.delta;
        if (delta?.reasoning) {
          yield { type: "reasoning.delta", text: delta.reasoning };
        }
        if (delta?.content) {
          yield { type: "content.delta", text: delta.content };
        }
        if (delta?.tool_calls) {
          for (const tc of delta.tool_calls) {
            const id = tc.id ?? `tool_${tc.index ?? 0}`;
            if (tc.function?.name && !openToolCalls.has(id)) {
              openToolCalls.add(id);
              yield { type: "tool.start", toolCallId: id, name: tc.function.name };
            }
            if (tc.function?.arguments) {
              yield { type: "tool.delta", toolCallId: id, argumentsDelta: tc.function.arguments };
            }
          }
        }
        if (chunk.usage) {
          yield {
            type: "usage",
            promptTokens: chunk.usage.prompt_tokens,
            completionTokens: chunk.usage.completion_tokens,
            reasoningTokens: chunk.usage.completion_tokens_details?.reasoning_tokens,
            cachedTokens: chunk.usage.prompt_tokens_details?.cached_tokens,
            totalTokens: chunk.usage.total_tokens,
            cost: chunk.usage.cost,
          };
        }
      }
    }
    yield { type: "generation.end" };
  } catch (e) {
    if (signal.aborted) {
      // Cancellation is not an error; the caller marks the generation cancelled.
      return;
    }
    const err = normalizeNetworkError(e, request.model);
    yield { type: "generation.error", message: err.message, retryable: err.retryable };
  } finally {
    reader.cancel().catch(() => {});
  }
}

/** The initial (and only) gateway implementation. */
export class OpenRouterGateway implements ModelGateway {
  async *stream(request: GenerationRequest): AsyncIterable<ChatStreamEvent> {
    yield* streamGeneration(request, crypto.randomUUID());
  }
  listModels(): Promise<ModelDefinition[]> {
    return fetchModels();
  }
}