import type OpenAI from "openai";
import type { ChatCompletionChunk, ChatCompletionMessageParam, ChatCompletionContentPart } from "openai/resources/chat/completions";
import type { UnifiedMessage, UnifiedStreamEvent, FinishReason, ProviderId } from "@/lib/ai/core/types";
import { inlineTextFile, isTextLike } from "@/lib/ai/core/content";
import { safeJson } from "@/lib/ai/core/stream-utils";

/**
 * OpenAI-shaped Chat Completions helpers, shared by xAI (and any future
 * OpenAI-compatible endpoint such as Groq, Together, local servers).
 */

export interface ChatCompletionsBuildOptions {
  /** Role used for the system prompt. */
  systemRole?: "system" | "developer";
  /** When the endpoint has no PDF support, text-like files are inlined; others are described. */
  inlineFiles?: boolean;
  /** Replay stored reasoning as `reasoning_content` on assistant turns (required by Kimi/DeepSeek thinking models when tools are used). */
  replayReasoningContent?: boolean;
  /** Provider-specific content part for PDFs (e.g. Mistral `document_url`). When absent, PDFs are described. */
  pdfPart?: (data: string, name: string) => ChatCompletionContentPart;
}

export function toChatCompletionMessages(system: string | undefined, messages: UnifiedMessage[], opts: ChatCompletionsBuildOptions = {}): ChatCompletionMessageParam[] {
  const out: ChatCompletionMessageParam[] = [];
  if (system?.trim()) out.push({ role: opts.systemRole ?? "system", content: system } as ChatCompletionMessageParam);
  for (const m of messages) {
    if (m.role === "system") {
      out.push({ role: opts.systemRole ?? "system", content: m.content.map((p) => (p.type === "text" ? p.text : "")).join("") } as ChatCompletionMessageParam);
      continue;
    }
    if (m.role === "tool") {
      for (const p of m.content) {
        if (p.type === "tool-result") out.push({ role: "tool", tool_call_id: p.toolCallId, content: typeof p.result === "string" ? p.result : JSON.stringify(p.result ?? null) });
      }
      continue;
    }
    if (m.role === "assistant") {
      const text = m.content.map((p) => (p.type === "text" ? p.text : "")).join("");
      const toolCalls = m.content.filter((p) => p.type === "tool-call");
      const msg: ChatCompletionMessageParam = { role: "assistant", content: text || null };
      if (opts.replayReasoningContent) {
        const reasoning = m.content.filter((p) => p.type === "reasoning").map((p) => (p.type === "reasoning" ? p.text : "")).join("");
        if (reasoning) (msg as unknown as Record<string, unknown>).reasoning_content = reasoning;
      }
      if (toolCalls.length) {
        msg.tool_calls = toolCalls.map((p) => (p.type === "tool-call" ? { id: p.id, type: "function" as const, function: { name: p.name, arguments: p.argumentsText ?? JSON.stringify(p.arguments) } } : null)).filter(Boolean) as NonNullable<typeof msg.tool_calls>;
      }
      // tool results embedded in the assistant unified message (single-turn tool loops)
      out.push(msg);
      for (const p of m.content) {
        if (p.type === "tool-result") out.push({ role: "tool", tool_call_id: p.toolCallId, content: typeof p.result === "string" ? p.result : JSON.stringify(p.result ?? null) });
      }
      continue;
    }
    // user
    const parts: ChatCompletionContentPart[] = [];
    for (const p of m.content) {
      if (p.type === "text") parts.push({ type: "text", text: p.text });
      else if (p.type === "image") parts.push({ type: "image_url", image_url: { url: `data:${p.mimeType};base64,${p.data}`, detail: "auto" } });
      else if (p.type === "file") {
        if (p.mimeType === "application/pdf" && opts.pdfPart) parts.push(opts.pdfPart(p.data, p.name));
        else if (isTextLike(p.mimeType, p.name) || opts.inlineFiles !== false) parts.push({ type: "text", text: isTextLike(p.mimeType, p.name) ? inlineTextFile(p.name, p.data) : `[Attached file "${p.name}" (${p.mimeType}) — this model cannot read binary files of this type.]` });
      }
    }
    if (parts.length === 1 && parts[0].type === "text") out.push({ role: "user", content: parts[0].text });
    else if (parts.length) out.push({ role: "user", content: parts });
  }
  return out;
}

export function mapChatFinish(reason: string | null | undefined): FinishReason {
  switch (reason) {
    case "stop":
      return "stop";
    case "length":
      return "length";
    case "tool_calls":
    case "function_call":
      return "tool-calls";
    case "content_filter":
      return "content-filter";
    default:
      return "other";
  }
}

interface ToolAcc {
  id: string;
  name: string;
  args: string;
  started: boolean;
}

/**
 * Normalize an OpenAI-shaped chat.completion.chunk stream into UnifiedStreamEvents.
 * Handles `delta.reasoning_content` (xAI, DeepSeek-style) and `delta.reasoning` variants.
 */
export async function* normalizeChatCompletionStream(stream: AsyncIterable<ChatCompletionChunk>, provider: ProviderId): AsyncIterable<UnifiedStreamEvent> {
  const tools = new Map<number, ToolAcc>();
  let finish: FinishReason = "other";
  let started = false;
  let sawUsage = false;
  let lastUpstream: string | undefined;
  for await (const chunk of stream) {
    if (!started) {
      started = true;
      yield { type: "start", id: chunk.id, model: chunk.model };
    }
    const upstream = (chunk as { provider?: string }).provider;
    if (upstream && upstream !== lastUpstream) {
      lastUpstream = upstream;
      yield { type: "provider-data", data: { upstreamProvider: upstream } };
    }
    const choice = chunk.choices?.[0];
    if (choice) {
      const annotations = (choice.delta as { annotations?: Array<{ type?: string; url_citation?: { url?: string; title?: string; content?: string } }> }).annotations;
      if (Array.isArray(annotations)) {
        for (const a of annotations) if (a.type === "url_citation" && a.url_citation?.url) yield { type: "citation", citation: { url: a.url_citation.url, title: a.url_citation.title, snippet: a.url_citation.content?.slice(0, 300), source: "web_search" } };
      }
      const delta = choice.delta as Omit<ChatCompletionChunk.Choice.Delta, "content"> & { content?: string | null | Array<{ type?: string; text?: string; thinking?: Array<{ type?: string; text?: string }> | string }>; reasoning_content?: string | null; reasoning?: string | null };
      const reasoning = delta.reasoning_content ?? delta.reasoning;
      if (reasoning) yield { type: "reasoning-delta", text: reasoning };
      if (typeof delta.content === "string" && delta.content) yield { type: "text-delta", text: delta.content };
      else if (Array.isArray(delta.content)) {
        // Mistral (Magistral) streams content as chunks: {type:"thinking", thinking:[{type:"text", text}]} | {type:"text", text}
        for (const c of delta.content) {
          if (c.type === "thinking") {
            const t = typeof c.thinking === "string" ? c.thinking : (c.thinking ?? []).map((x) => x.text ?? "").join("");
            if (t) yield { type: "reasoning-delta", text: t };
          } else if (c.text) yield { type: "text-delta", text: c.text };
        }
      }
      if (delta.tool_calls) {
        for (const tc of delta.tool_calls) {
          const idx = tc.index ?? 0;
          let acc = tools.get(idx);
          if (!acc) {
            acc = { id: tc.id ?? `call_${idx}`, name: tc.function?.name ?? "", args: "", started: false };
            tools.set(idx, acc);
          }
          if (tc.id) acc.id = tc.id;
          if (tc.function?.name) acc.name = tc.function.name;
          if (!acc.started && acc.name) {
            acc.started = true;
            yield { type: "tool-start", id: acc.id, name: acc.name };
          }
          if (tc.function?.arguments) {
            acc.args += tc.function.arguments;
            yield { type: "tool-delta", id: acc.id, argumentsDelta: tc.function.arguments };
          }
        }
      }
      if (choice.finish_reason) finish = mapChatFinish(choice.finish_reason);
    }
    const usage = chunk.usage as (ChatCompletionChunk["usage"] & { prompt_tokens_details?: { cached_tokens?: number | null } | null; completion_tokens_details?: { reasoning_tokens?: number | null } | null; prompt_cache_hit_tokens?: number | null; cost?: number | null; cached_tokens?: number | null }) | null | undefined;
    if (usage) {
      sawUsage = true;
      const reasoningTokens = usage.completion_tokens_details?.reasoning_tokens ?? undefined;
      const cached = usage.prompt_tokens_details?.cached_tokens ?? usage.prompt_cache_hit_tokens ?? usage.cached_tokens ?? undefined;
      // xAI: total_tokens = prompt + completion + reasoning; completion excludes reasoning. We bill reasoning as output.
      const outputTokens = (usage.completion_tokens ?? 0) + (provider === "xai" ? reasoningTokens ?? 0 : 0);
      yield {
        type: "usage",
        usage: {
          inputTokens: usage.prompt_tokens ?? 0,
          outputTokens,
          cachedInputTokens: cached ?? undefined,
          reasoningTokens,
          totalTokens: usage.total_tokens ?? undefined,
        },
      };
      const cost = (usage as { cost_in_usd_ticks?: number }).cost_in_usd_ticks;
      if (typeof cost === "number") yield { type: "provider-data", data: { exactCostUsd: cost / 1e10 } };
      if (typeof usage.cost === "number") yield { type: "provider-data", data: { exactCostUsd: usage.cost } }; // OpenRouter usage accounting
    }
  }
  for (const acc of tools.values()) {
    if (!acc.started) yield { type: "tool-start", id: acc.id, name: acc.name };
    yield { type: "tool-end", id: acc.id, name: acc.name, arguments: safeJson(acc.args), argumentsText: acc.args };
  }
  if (!sawUsage) yield { type: "usage", usage: { inputTokens: 0, outputTokens: 0 } };
  yield { type: "finish", reason: finish };
}

export type OpenAIClient = OpenAI;