import OpenAI from "openai";
import type { ChatCompletionCreateParamsStreaming, ChatCompletionTool } from "openai/resources/chat/completions";
import type { ResponseCreateParamsStreaming, Tool as ResponsesTool } from "openai/resources/responses/responses";
import {
  type AIProviderAdapter,
  type PolyModel,
  type UnifiedChatRequest,
  type UnifiedChatResponse,
  type UnifiedStreamEvent,
  type ValidationResult,
  type TokenEstimate,
  PolyProviderError,
  modelKey,
} from "@/lib/ai/core/types";
import { normalizeGenericError, refineByMessage, codeFromStatus, isRetryableStatus, parseRetryAfter } from "@/lib/ai/core/errors";
import { filterSettings, heuristicTokens } from "@/lib/ai/core/normalize";
import { collectStream } from "@/lib/ai/core/stream-utils";
import { toChatCompletionMessages, normalizeChatCompletionStream } from "../shared/openai-compat/chat-completions";
import { toResponsesInput, normalizeResponsesStream } from "../shared/openai-compat/responses";
import { XAI_CATALOG } from "./catalog";

const BASE_URL = "https://api.x.ai/v1";
const DEFAULT_TIMEOUT_MS = 10 * 60_000;

function client(apiKey: string, timeoutMs = DEFAULT_TIMEOUT_MS) {
  return new OpenAI({ apiKey, baseURL: BASE_URL, maxRetries: 1, timeout: timeoutMs });
}

interface XaiLanguageModel {
  id: string;
  input_modalities?: string[];
  output_modalities?: string[];
  prompt_text_token_price?: number;
  cached_prompt_text_token_price?: number;
  completion_text_token_price?: number;
  prompt_text_token_price_long_context?: number;
  cached_prompt_text_token_price_long_context?: number;
  completion_text_token_price_long_context?: number;
  long_context_threshold?: number;
  aliases?: string[];
  created?: number;
}

/** xAI prices are cents per 100M tokens → $/M = value / 10 000 (verified against the pricing page). */
const perMillion = (v?: number) => (typeof v === "number" ? v / 10_000 : undefined);

export function normalizeXaiModel(m: XaiLanguageModel): PolyModel {
  const cat = XAI_CATALOG.get(m.id);
  const vision = m.input_modalities?.includes("image") ?? true;
  return {
    key: modelKey("xai", m.id),
    id: m.id,
    provider: "xai",
    displayName: cat?.displayName ?? prettify(m.id),
    family: cat?.family ?? m.id.split("-").slice(0, 2).join(" "),
    capabilities: cat?.capabilities ?? {
      text: true,
      vision,
      audioInput: false,
      audioOutput: false,
      imageGeneration: false,
      video: false,
      reasoning: !m.id.includes("non-reasoning"),
      tools: !m.id.includes("multi-agent"),
      structuredOutput: true,
      streaming: true,
      files: false,
      webSearch: true,
    },
    limits: cat?.limits ?? {},
    parameters: cat?.parameters ?? { temperature: true, topP: true, topK: true, maxTokens: true, seed: true, stop: false, frequencyPenalty: false, presencePenalty: false, reasoningEffort: false },
    status: cat?.status ?? "unknown",
    pricing: {
      inputPerMillion: perMillion(m.prompt_text_token_price),
      cachedInputPerMillion: perMillion(m.cached_prompt_text_token_price),
      outputPerMillion: perMillion(m.completion_text_token_price),
      longContext: m.long_context_threshold
        ? {
            thresholdTokens: m.long_context_threshold,
            inputPerMillion: perMillion(m.prompt_text_token_price_long_context),
            cachedInputPerMillion: perMillion(m.cached_prompt_text_token_price_long_context),
            outputPerMillion: perMillion(m.completion_text_token_price_long_context),
          }
        : undefined,
      source: "xai:/v1/language-models",
      asOf: new Date().toISOString().slice(0, 10),
    },
    metadata: { ...(cat?.metadata ?? {}), aliases: m.aliases ?? [], sortWeight: cat?.sortWeight ?? 0, createdAt: m.created ? new Date(m.created * 1000).toISOString() : undefined },
  };
}

function prettify(id: string): string {
  return id.replace(/^grok/, "Grok").replace(/-/g, " ").replace(/\b(\w)/g, (c) => c.toUpperCase());
}

async function fetchLanguageModels(apiKey: string, signal?: AbortSignal): Promise<XaiLanguageModel[]> {
  const res = await fetch(`${BASE_URL}/language-models`, { headers: { Authorization: `Bearer ${apiKey}` }, signal });
  if (!res.ok) {
    const body = (await res.json().catch(() => ({}))) as { code?: string; error?: string };
    throw Object.assign(new Error(body.error ?? `HTTP ${res.status}`), { status: res.status, error: body, headers: res.headers });
  }
  const data = (await res.json()) as { models?: XaiLanguageModel[] };
  return data.models ?? [];
}

function buildChatParams(req: UnifiedChatRequest): ChatCompletionCreateParamsStreaming {
  const { settings } = filterSettings(req.settings, req.modelInfo);
  const params: ChatCompletionCreateParamsStreaming = {
    model: req.model,
    messages: toChatCompletionMessages(req.system, req.messages, { systemRole: "system", inlineFiles: true }),
    stream: true,
    stream_options: { include_usage: true },
  };
  if (settings.maxTokens !== undefined) params.max_completion_tokens = settings.maxTokens;
  if (settings.temperature !== undefined) params.temperature = settings.temperature;
  if (settings.topP !== undefined) params.top_p = settings.topP;
  if (settings.seed !== undefined) params.seed = settings.seed;
  if (settings.stop?.length) params.stop = settings.stop.slice(0, 4);
  if (settings.topK !== undefined) (params as unknown as Record<string, unknown>).top_k = settings.topK;
  if (settings.reasoningEffort) {
    const effort = settings.reasoningEffort === "minimal" ? "low" : settings.reasoningEffort === "max" ? "xhigh" : settings.reasoningEffort;
    (params as unknown as Record<string, unknown>).reasoning_effort = effort;
  }
  if (settings.responseFormat?.type === "json_schema" && settings.responseFormat.schema) {
    params.response_format = { type: "json_schema", json_schema: { name: settings.responseFormat.schemaName ?? "response", schema: settings.responseFormat.schema, strict: settings.responseFormat.strict ?? true } };
  } else if (settings.responseFormat?.type === "json") {
    params.response_format = { type: "json_object" };
  }
  if (req.tools?.length) {
    params.tools = req.tools.map<ChatCompletionTool>((t) => ({ type: "function", function: { name: t.name, description: t.description, parameters: t.parameters, strict: t.strict } }));
    const tc = settings.toolChoice;
    if (tc === "none" || tc === "required" || tc === "auto") params.tool_choice = tc;
    else if (tc && typeof tc === "object") params.tool_choice = { type: "function", function: { name: tc.name } };
  }
  return params;
}

function buildResponsesParams(req: UnifiedChatRequest): ResponseCreateParamsStreaming {
  const { settings } = filterSettings(req.settings, req.modelInfo);
  const params: ResponseCreateParamsStreaming = {
    model: req.model,
    input: toResponsesInput(req.messages, { nativeFiles: false, replayReasoning: false }),
    stream: true,
    store: false,
  };
  if (req.system?.trim()) params.instructions = req.system;
  if (settings.maxTokens !== undefined) params.max_output_tokens = settings.maxTokens;
  if (settings.temperature !== undefined) params.temperature = settings.temperature;
  if (settings.topP !== undefined) params.top_p = settings.topP;
  if (settings.reasoningEffort) {
    const effort = settings.reasoningEffort === "minimal" ? "low" : settings.reasoningEffort === "max" ? "xhigh" : settings.reasoningEffort;
    params.reasoning = { effort: effort as "low" | "medium" | "high", summary: "auto" };
  } else if (req.modelInfo?.capabilities.reasoning && settings.includeReasoning !== false) {
    params.reasoning = { summary: "auto" };
  }
  if (settings.responseFormat?.type === "json_schema" && settings.responseFormat.schema) {
    params.text = { format: { type: "json_schema", name: settings.responseFormat.schemaName ?? "response", schema: settings.responseFormat.schema, strict: settings.responseFormat.strict ?? true } };
  }
  const tools: ResponsesTool[] = [];
  if (settings.webSearch) tools.push({ type: "web_search" } as ResponsesTool);
  for (const t of req.tools ?? []) tools.push({ type: "function", name: t.name, description: t.description, parameters: t.parameters, strict: t.strict ?? true });
  if (tools.length) params.tools = tools;
  return params;
}

export const xaiAdapter: AIProviderAdapter = {
  id: "xai",
  name: "xAI",
  keyDocsUrl: "https://console.x.ai",
  keyPrefixHint: "xai-",

  async validateApiKey(apiKey, signal): Promise<ValidationResult> {
    const t0 = Date.now();
    try {
      const models = await fetchLanguageModels(apiKey, signal);
      return { ok: true, modelsAvailable: models.length, latencyMs: Date.now() - t0 };
    } catch (e) {
      return { ok: false, error: this.normalizeError(e).toJSON(), latencyMs: Date.now() - t0 };
    }
  },

  async listModels(apiKey, signal): Promise<PolyModel[]> {
    try {
      const models = await fetchLanguageModels(apiKey, signal);
      return models.filter((m) => (m.output_modalities ?? ["text"]).includes("text")).map(normalizeXaiModel);
    } catch (e) {
      throw this.normalizeError(e);
    }
  },

  async chat(req): Promise<UnifiedChatResponse> {
    return collectStream("xai", req.model, this.streamChat(req));
  },

  async *streamChat(req: UnifiedChatRequest): AsyncIterable<UnifiedStreamEvent> {
    const meta = (req.modelInfo?.metadata ?? {}) as Record<string, unknown>;
    const useResponses = Boolean(req.settings?.webSearch) || meta.responsesOnly === true;
    const c = client(req.apiKey, req.timeoutMs);
    try {
      if (useResponses) {
        const stream = await c.responses.create(buildResponsesParams(req), { signal: req.signal });
        yield* normalizeResponsesStream(stream, "xai");
      } else {
        const stream = await c.chat.completions.create(buildChatParams(req), { signal: req.signal });
        yield* normalizeChatCompletionStream(stream, "xai");
      }
    } catch (e) {
      yield { type: "error", error: this.normalizeError(e).toJSON() };
    }
  },

  async estimateTokens(req): Promise<TokenEstimate> {
    const text = req.messages.map((m) => m.content.map((p) => (p.type === "text" ? p.text : "")).join(" ")).join(" ") + (req.system ?? "");
    return { inputTokens: heuristicTokens(text), method: "heuristic" };
  },

  normalizeError(error: unknown): PolyProviderError {
    if (error instanceof PolyProviderError) return error;
    const e = error as { status?: number; error?: { code?: string; error?: string; message?: string } | string; message?: string; headers?: Headers };
    const body = typeof e?.error === "object" ? e.error : undefined;
    const message = body?.error ?? body?.message ?? e?.message ?? String(error);
    const status = e?.status;
    if (status !== undefined || body) {
      let code = codeFromStatus(status);
      // xAI returns HTTP 400 for bad keys and for unknown models — disambiguate by text.
      if (/incorrect api key|bad credentials|unauthenticated|no credentials/i.test(message) || body?.code?.startsWith("unauthenticated")) code = "INVALID_API_KEY";
      else if (/model not found/i.test(message)) code = "MODEL_NOT_FOUND";
      else if (/does not support parameter|invalid-argument|invalid_image/i.test(message + (body?.code ?? ""))) code = refineByMessage("INVALID_PARAMETER", message);
      else code = refineByMessage(code, message);
      if (body?.code === "invalid_image") code = "INVALID_PARAMETER";
      return new PolyProviderError({
        code,
        message: code === "INVALID_API_KEY" ? "Invalid API key" : message.slice(0, 600),
        provider: "xai",
        status,
        retryable: isRetryableStatus(status) && code !== "INVALID_API_KEY",
        retryAfterMs: parseRetryAfter(e?.headers ?? null),
        providerCode: body?.code,
        cause: error,
      });
    }
    return normalizeGenericError("xai", error);
  },
};