import type { ChatModel, ModelProvider } from "./modelRegistry.js";
import { apiKeyForProvider } from "./modelRegistry.js";
import { normalizeUsage, type Usage } from "./tokens.js";

export type ChatRequest = {
  model: ChatModel;
  provider: ModelProvider;
  systemPrompt?: string;
  prompt: string;
  temperature: number;
  maxTokens: number;
};

export type ChatResponse = {
  output: string;
  reasoning?: string;
  usage: Usage;
  rawModel: string;
};

export type StreamDelta = {
  channel: "reasoning" | "output";
  text: string;
};

export type ChatStreamRequest = ChatRequest & {
  onDelta: (delta: StreamDelta) => void;
  signal?: AbortSignal;
};

type OpenAIResponse = {
  choices?: Array<{ message?: { content?: string }; text?: string }>;
  usage?: {
    prompt_tokens?: number;
    completion_tokens?: number;
    total_tokens?: number;
    input_tokens?: number;
    output_tokens?: number;
  };
};

type OpenAIStreamDelta = {
  content?: string;
  reasoning?: string;
  reasoning_content?: string;
  thinking?: string;
  reasoning_details?: Array<{ text?: string; content?: string }>;
};

type OpenAIStreamChunk = {
  choices?: Array<{
    delta?: OpenAIStreamDelta;
  }>;
  usage?: OpenAIResponse["usage"];
};

type AnthropicResponse = {
  content?: Array<{ type?: string; text?: string }>;
  usage?: {
    input_tokens?: number;
    output_tokens?: number;
  };
};

type AnthropicStreamEvent = {
  type?: string;
  content_block?: {
    type?: string;
    text?: string;
    thinking?: string;
  };
  delta?: {
    type?: string;
    text?: string;
    thinking?: string;
  };
  message?: {
    usage?: {
      input_tokens?: number;
      output_tokens?: number;
    };
  };
  usage?: {
    input_tokens?: number;
    output_tokens?: number;
  };
};

type GeminiResponse = {
  candidates?: Array<{
    content?: {
      parts?: Array<{ text?: string }>;
    };
  }>;
  usageMetadata?: {
    promptTokenCount?: number;
    candidatesTokenCount?: number;
    totalTokenCount?: number;
  };
};

type GeminiStreamChunk = GeminiResponse;

const readErrorBody = async (response: Response) => {
  const text = await response.text();
  if (!text) {
    return response.statusText;
  }

  try {
    const json = JSON.parse(text) as { error?: { message?: string }; message?: string };
    return json.error?.message ?? json.message ?? text;
  } catch {
    return text;
  }
};

const postJson = async <T>(
  url: string,
  headers: Record<string, string>,
  body: unknown
): Promise<T> => {
  const response = await fetch(url, {
    method: "POST",
    headers: {
      "content-type": "application/json",
      ...headers
    },
    body: JSON.stringify(body)
  });

  if (!response.ok) {
    throw new Error(await readErrorBody(response));
  }

  return (await response.json()) as T;
};

const postJsonStream = async (
  url: string,
  headers: Record<string, string>,
  body: unknown,
  signal?: AbortSignal
) => {
  const response = await fetch(url, {
    method: "POST",
    headers: {
      "content-type": "application/json",
      ...headers
    },
    body: JSON.stringify(body),
    signal
  });

  if (!response.ok) {
    throw new Error(await readErrorBody(response));
  }

  if (!response.body) {
    throw new Error("Provider returned an empty stream");
  }

  return response;
};

const readSse = async (
  response: Response,
  onEvent: (event: { event?: string; data: string }) => void
) => {
  const reader = response.body?.getReader();
  if (!reader) {
    throw new Error("Provider stream is not readable");
  }

  const decoder = new TextDecoder();
  let buffer = "";
  let eventName = "";
  let dataLines: string[] = [];

  const flush = () => {
    if (!dataLines.length) {
      eventName = "";
      return;
    }

    onEvent({
      event: eventName || undefined,
      data: dataLines.join("\n")
    });
    eventName = "";
    dataLines = [];
  };

  const handleLine = (rawLine: string) => {
    const line = rawLine.endsWith("\r") ? rawLine.slice(0, -1) : rawLine;
    if (!line) {
      flush();
      return;
    }

    if (line.startsWith("event:")) {
      eventName = line.slice(6).trim();
      return;
    }

    if (line.startsWith("data:")) {
      dataLines.push(line.slice(5).trimStart());
      return;
    }

    if (line.trim().startsWith("{")) {
      onEvent({ data: line.trim() });
    }
  };

  while (true) {
    const { value, done } = await reader.read();
    if (done) {
      break;
    }

    buffer += decoder.decode(value, { stream: true });
    const lines = buffer.split("\n");
    buffer = lines.pop() ?? "";
    for (const line of lines) {
      handleLine(line);
    }
  }

  buffer += decoder.decode();
  if (buffer) {
    handleLine(buffer);
  }
  flush();
};

const endpoint = (provider: ModelProvider, path: string) => {
  const baseUrl = provider.baseUrl?.replace(/\/$/, "");
  if (!baseUrl) {
    throw new Error(`${provider.name} base URL is not configured`);
  }

  return `${baseUrl}${path}`;
};

const composePrompt = (systemPrompt: string | undefined, prompt: string) =>
  [systemPrompt, prompt].filter(Boolean).join("\n\n");

const openAICompatibleMessages = (systemPrompt: string | undefined, prompt: string) => [
  ...(systemPrompt ? [{ role: "system", content: systemPrompt }] : []),
  { role: "user", content: prompt }
];

export const callOpenAICompatible = async ({
  provider,
  model,
  systemPrompt,
  prompt,
  temperature,
  maxTokens
}: ChatRequest): Promise<ChatResponse> => {
  const apiKey = apiKeyForProvider(provider);
  if (!apiKey) {
    throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
  }

  const json = await postJson<OpenAIResponse>(
    endpoint(provider, "/chat/completions"),
    { authorization: `Bearer ${apiKey}` },
    {
      model: model.model,
      messages: openAICompatibleMessages(systemPrompt, prompt),
      temperature,
      max_tokens: maxTokens
    }
  );

  const output = json.choices?.[0]?.message?.content ?? json.choices?.[0]?.text ?? "";
  const usage = normalizeUsage(composePrompt(systemPrompt, prompt), output, {
    inputTokens: json.usage?.prompt_tokens ?? json.usage?.input_tokens,
    outputTokens: json.usage?.completion_tokens ?? json.usage?.output_tokens,
    totalTokens: json.usage?.total_tokens,
    estimated: !json.usage
  });

  return { output, usage, rawModel: model.model };
};

const reasoningTextFromDelta = (delta?: OpenAIStreamDelta) => {
  const details = delta?.reasoning_details
    ?.map((detail) => detail.text ?? detail.content ?? "")
    .join("");
  return delta?.reasoning_content ?? delta?.reasoning ?? delta?.thinking ?? details ?? "";
};

export const streamOpenAICompatible = async ({
  provider,
  model,
  systemPrompt,
  prompt,
  temperature,
  maxTokens,
  onDelta,
  signal
}: ChatStreamRequest): Promise<ChatResponse> => {
  const apiKey = apiKeyForProvider(provider);
  if (!apiKey) {
    throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
  }

  let output = "";
  let reasoning = "";
  let usage: OpenAIResponse["usage"] | undefined;

  const response = await postJsonStream(
    endpoint(provider, "/chat/completions"),
    { authorization: `Bearer ${apiKey}` },
    {
      model: model.model,
      messages: openAICompatibleMessages(systemPrompt, prompt),
      temperature,
      max_tokens: maxTokens,
      stream: true,
      stream_options: { include_usage: true }
    },
    signal
  );

  await readSse(response, ({ data }) => {
    if (data === "[DONE]") {
      return;
    }

    const chunk = JSON.parse(data) as OpenAIStreamChunk;
    usage = chunk.usage ?? usage;
    const delta = chunk.choices?.[0]?.delta;
    const reasoningDelta = reasoningTextFromDelta(delta);
    const outputDelta = delta?.content ?? "";

    if (reasoningDelta) {
      reasoning += reasoningDelta;
      onDelta({ channel: "reasoning", text: reasoningDelta });
    }

    if (outputDelta) {
      output += outputDelta;
      onDelta({ channel: "output", text: outputDelta });
    }
  });

  return {
    output,
    reasoning,
    usage: normalizeUsage(composePrompt(systemPrompt, prompt), output, {
      inputTokens: usage?.prompt_tokens ?? usage?.input_tokens,
      outputTokens: usage?.completion_tokens ?? usage?.output_tokens,
      totalTokens: usage?.total_tokens,
      estimated: !usage
    }),
    rawModel: model.model
  };
};

export const callAnthropic = async ({
  provider,
  model,
  systemPrompt,
  prompt,
  temperature,
  maxTokens
}: ChatRequest): Promise<ChatResponse> => {
  const apiKey = apiKeyForProvider(provider);
  if (!apiKey) {
    throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
  }

  const json = await postJson<AnthropicResponse>(
    endpoint(provider, "/messages"),
    {
      "x-api-key": apiKey,
      "anthropic-version": "2023-06-01"
    },
    {
      model: model.model,
      system: systemPrompt || undefined,
      messages: [{ role: "user", content: prompt }],
      temperature,
      max_tokens: maxTokens
    }
  );

  const output = json.content?.map((part) => part.text ?? "").join("\n").trim() ?? "";
  const usage = normalizeUsage(composePrompt(systemPrompt, prompt), output, {
    inputTokens: json.usage?.input_tokens,
    outputTokens: json.usage?.output_tokens,
    estimated: !json.usage
  });

  return { output, usage, rawModel: model.model };
};

export const streamAnthropic = async ({
  provider,
  model,
  systemPrompt,
  prompt,
  temperature,
  maxTokens,
  onDelta,
  signal
}: ChatStreamRequest): Promise<ChatResponse> => {
  const apiKey = apiKeyForProvider(provider);
  if (!apiKey) {
    throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
  }

  let output = "";
  let reasoning = "";
  let inputTokens: number | undefined;
  let outputTokens: number | undefined;

  const response = await postJsonStream(
    endpoint(provider, "/messages"),
    {
      "x-api-key": apiKey,
      "anthropic-version": "2023-06-01"
    },
    {
      model: model.model,
      system: systemPrompt || undefined,
      messages: [{ role: "user", content: prompt }],
      temperature,
      max_tokens: maxTokens,
      stream: true
    },
    signal
  );

  await readSse(response, ({ data }) => {
    const event = JSON.parse(data) as AnthropicStreamEvent;
    inputTokens = event.message?.usage?.input_tokens ?? event.usage?.input_tokens ?? inputTokens;
    outputTokens = event.message?.usage?.output_tokens ?? event.usage?.output_tokens ?? outputTokens;

    const startText =
      event.type === "content_block_start" && event.content_block?.type === "text"
        ? event.content_block.text ?? ""
        : "";
    const startThinking =
      event.type === "content_block_start" && event.content_block?.type === "thinking"
        ? event.content_block.thinking ?? ""
        : "";
    const outputDelta = event.delta?.type === "text_delta" ? event.delta.text ?? "" : "";
    const reasoningDelta =
      event.delta?.type === "thinking_delta" ? event.delta.thinking ?? "" : "";

    if (startThinking || reasoningDelta) {
      const text = startThinking || reasoningDelta;
      reasoning += text;
      onDelta({ channel: "reasoning", text });
    }

    if (startText || outputDelta) {
      const text = startText || outputDelta;
      output += text;
      onDelta({ channel: "output", text });
    }
  });

  return {
    output,
    reasoning,
    usage: normalizeUsage(composePrompt(systemPrompt, prompt), output, {
      inputTokens,
      outputTokens,
      estimated: inputTokens === undefined && outputTokens === undefined
    }),
    rawModel: model.model
  };
};

export const callGemini = async ({
  provider,
  model,
  systemPrompt,
  prompt,
  temperature,
  maxTokens
}: ChatRequest): Promise<ChatResponse> => {
  const apiKey = apiKeyForProvider(provider);
  if (!apiKey) {
    throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
  }

  const baseUrl = provider.baseUrl?.replace(/\/$/, "");
  const url = `${baseUrl}/models/${encodeURIComponent(model.model)}:generateContent?key=${encodeURIComponent(apiKey)}`;
  const json = await postJson<GeminiResponse>(
    url,
    {},
    {
      systemInstruction: systemPrompt ? { parts: [{ text: systemPrompt }] } : undefined,
      contents: [{ role: "user", parts: [{ text: prompt }] }],
      generationConfig: {
        temperature,
        maxOutputTokens: maxTokens
      }
    }
  );

  const output =
    json.candidates?.[0]?.content?.parts?.map((part) => part.text ?? "").join("\n").trim() ?? "";
  const usage = normalizeUsage(composePrompt(systemPrompt, prompt), output, {
    inputTokens: json.usageMetadata?.promptTokenCount,
    outputTokens: json.usageMetadata?.candidatesTokenCount,
    totalTokens: json.usageMetadata?.totalTokenCount,
    estimated: !json.usageMetadata
  });

  return { output, usage, rawModel: model.model };
};

export const streamGemini = async ({
  provider,
  model,
  systemPrompt,
  prompt,
  temperature,
  maxTokens,
  onDelta,
  signal
}: ChatStreamRequest): Promise<ChatResponse> => {
  const apiKey = apiKeyForProvider(provider);
  if (!apiKey) {
    throw new Error(`${provider.name} API key is not configured (${provider.apiKeyEnv})`);
  }

  const baseUrl = provider.baseUrl?.replace(/\/$/, "");
  const url = `${baseUrl}/models/${encodeURIComponent(model.model)}:streamGenerateContent?alt=sse&key=${encodeURIComponent(apiKey)}`;
  let output = "";
  let reasoning = "";
  let usage: GeminiResponse["usageMetadata"] | undefined;

  const response = await postJsonStream(
    url,
    {},
    {
      systemInstruction: systemPrompt ? { parts: [{ text: systemPrompt }] } : undefined,
      contents: [{ role: "user", parts: [{ text: prompt }] }],
      generationConfig: {
        temperature,
        maxOutputTokens: maxTokens
      }
    },
    signal
  );

  await readSse(response, ({ data }) => {
    const chunk = JSON.parse(data) as GeminiStreamChunk;
    usage = chunk.usageMetadata ?? usage;
    const parts = chunk.candidates?.[0]?.content?.parts ?? [];

    for (const part of parts) {
      const text = part.text ?? "";
      if (!text) {
        continue;
      }

      if ("thought" in part && part.thought === true) {
        reasoning += text;
        onDelta({ channel: "reasoning", text });
      } else {
        output += text;
        onDelta({ channel: "output", text });
      }
    }
  });

  return {
    output,
    reasoning,
    usage: normalizeUsage(composePrompt(systemPrompt, prompt), output, {
      inputTokens: usage?.promptTokenCount,
      outputTokens: usage?.candidatesTokenCount,
      totalTokens: usage?.totalTokenCount,
      estimated: !usage
    }),
    rawModel: model.model
  };
};

export const callProvider = async (request: ChatRequest): Promise<ChatResponse> => {
  switch (request.provider.kind) {
    case "openai-compatible":
      return callOpenAICompatible(request);
    case "anthropic":
      return callAnthropic(request);
    case "gemini":
      return callGemini(request);
    default: {
      const provider: never = request.provider.kind;
      throw new Error(`Unsupported provider kind: ${provider}`);
    }
  }
};

export const streamProvider = async (request: ChatStreamRequest): Promise<ChatResponse> => {
  switch (request.provider.kind) {
    case "openai-compatible":
      return streamOpenAICompatible(request);
    case "anthropic":
      return streamAnthropic(request);
    case "gemini":
      return streamGemini(request);
    default: {
      const provider: never = request.provider.kind;
      throw new Error(`Unsupported provider kind: ${provider}`);
    }
  }
};