import type { GenerateChunk, GenerateOpts, LLMProvider } from '../types';
import { sseLines } from '../sse';

/**
 * Generic chat-completions streamer used by both OpenAI and the
 * "openai-compatible" provider (OpenRouter, LiteLLM, vLLM, Together,
 * etc.). The wire format is the same — only the base URL differs.
 *
 * Streaming:
 *   - data: {choices:[{delta:{content:"..."}}]}
 *   - data: {choices:[{delta:{}}]}                  ← end-of-content
 *   - data: {usage:{prompt_tokens, completion_tokens}}  (when stream_options.include_usage)
 *   - data: [DONE]
 */
export async function* generateOpenAIStyle(
  opts: GenerateOpts,
  baseUrl: string,
  providerLabel: string,
): AsyncGenerator<GenerateChunk, void, void> {
  if (!opts.apiKey) {
    yield { type: 'error', message: `${providerLabel} API key is required.` };
    return;
  }
  const url = baseUrl.replace(/\/$/, '') + '/chat/completions';
  let res: Response;
  try {
    res = await fetch(url, {
      method: 'POST',
      headers: {
        'content-type': 'application/json',
        authorization: `Bearer ${opts.apiKey}`,
      },
      body: JSON.stringify({
        model: opts.model,
        stream: true,
        stream_options: { include_usage: true },
        ...(opts.maxOutputTokens != null
          ? { max_completion_tokens: opts.maxOutputTokens }
          : {}),
        messages: [
          { role: 'system', content: opts.systemPrompt },
          { role: 'user', content: opts.userPrompt },
        ],
      }),
      signal: opts.signal,
    });
  } catch (e) {
    yield {
      type: 'error',
      message: `${providerLabel} unreachable: ${(e as Error).message}`,
    };
    return;
  }
  if (!res.ok) {
    yield {
      type: 'error',
      message: `${providerLabel} HTTP ${res.status}: ${(await res.text().catch(() => '')).slice(0, 500)}`,
    };
    return;
  }
  let tokensIn: number | undefined;
  let tokensOut: number | undefined;
  try {
    for await (const data of sseLines(res)) {
      if (data === '[DONE]') break;
      let evt: any;
      try {
        evt = JSON.parse(data);
      } catch {
        continue;
      }
      const delta = evt.choices?.[0]?.delta?.content;
      if (delta) yield { type: 'text', delta };
      if (evt.usage) {
        tokensIn = evt.usage.prompt_tokens;
        tokensOut = evt.usage.completion_tokens;
      }
    }
    yield { type: 'usage', tokensIn, tokensOut };
    yield { type: 'done' };
  } catch (e) {
    yield {
      type: 'error',
      message: `${providerLabel} stream error: ${(e as Error).message}`,
    };
  }
}

export const openai: LLMProvider = {
  id: 'openai',
  label: 'OpenAI',
  requiresApiKey: true,
  requiresBaseUrl: false,
  generate(opts) {
    return generateOpenAIStyle(opts, 'https://api.openai.com/v1', 'OpenAI');
  },
};

export const openaiCompatible: LLMProvider = {
  id: 'openai-compatible',
  label: 'OpenAI-compatible (custom URL)',
  requiresApiKey: true,
  requiresBaseUrl: true,
  async *generate(opts) {
    if (!opts.baseUrl) {
      yield {
        type: 'error',
        message:
          'Base URL is required (e.g. https://openrouter.ai/api/v1, your LiteLLM gateway, etc.).',
      };
      return;
    }
    yield* generateOpenAIStyle(opts, opts.baseUrl, 'OpenAI-compatible');
  },
};
