import type { Api, Model } from "@earendil-works/pi-ai";
import { openrouterProvider } from "@earendil-works/pi-ai/providers/openrouter";
import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";

const CATALOG_URL = "https://openrouter.ai/api/v1/models";
const REFRESH_INTERVAL_MS = 4 * 60 * 60 * 1000;
const CACHE_MARKER = "openrouter-direct";

type OpenRouterModel = {
  id: string;
  name: string;
  context_length?: number;
  architecture?: { input_modalities?: string[]; output_modalities?: string[]; modality?: string };
  pricing?: { prompt?: string; completion?: string; input_cache_read?: string; input_cache_write?: string };
  top_provider?: { context_length?: number; max_completion_tokens?: number };
  supported_parameters?: string[];
  reasoning?: { mandatory?: boolean; supported_efforts?: string[] };
};

function catalogModel(entry: OpenRouterModel): Model<Api> | undefined {
  if (!entry.id || !entry.name) return;
  if (entry.architecture?.output_modalities && !entry.architecture.output_modalities.includes("text")) return;
  if (entry.architecture?.input_modalities && !entry.architecture.input_modalities.includes("text")) return;

  const rate = (value: string | undefined): number => {
    const cost = Number(value ?? 0) * 1_000_000;
    return Number.isFinite(cost) && cost >= 0 ? Number(cost.toFixed(6)) : 0;
  };
  const efforts = entry.reasoning?.supported_efforts;
  const thinkingLevelMap = efforts?.length
    ? {
        off: entry.reasoning?.mandatory ? null : "none",
        ...Object.fromEntries(
          (["minimal", "low", "medium", "high", "xhigh", "max"] as const).map((level) => [
            level,
            efforts.includes(level) ? level : null,
          ]),
        ),
      }
    : entry.reasoning?.mandatory ? { off: null } : undefined;
  const anthropic = entry.id.startsWith("anthropic/") && !entry.id.endsWith(":batch");
  return {
    id: entry.id,
    name: entry.name,
    provider: "openrouter",
    api: anthropic ? "anthropic-messages" : "openai-completions",
    baseUrl: anthropic ? "https://openrouter.ai/api" : "https://openrouter.ai/api/v1",
    reasoning: entry.supported_parameters?.includes("reasoning") ?? false,
    ...(thinkingLevelMap && { thinkingLevelMap }),
    input: entry.architecture?.input_modalities?.includes("image") || entry.architecture?.modality?.includes("image")
      ? ["text", "image"]
      : ["text"],
    cost: {
      input: rate(entry.pricing?.prompt),
      output: rate(entry.pricing?.completion),
      cacheRead: rate(entry.pricing?.input_cache_read),
      cacheWrite: rate(entry.pricing?.input_cache_write),
    },
    contextWindow: entry.top_provider?.context_length || entry.context_length || 4096,
    maxTokens: entry.top_provider?.max_completion_tokens || 4096,
  };
}

export default function (pi: ExtensionAPI) {
  const builtIn = openrouterProvider();
  let liveModels: readonly Model<Api>[] = [];

  pi.registerProvider({
    ...builtIn,
    getModels: () => {
      // Keep Pi's hand-tuned settings for known models; add models missing from its catalog.
      const models = new Map<string, Model<Api>>(liveModels.map((model) => [model.id, model]));
      for (const model of builtIn.getModels()) models.set(model.id, model);
      return [...models.values()].filter((model) => !/^(openai|anthropic)\//.test(model.id));
    },
    async refreshModels(context) {
      const stored = context.stored;
      if (stored?.models.length) {
        await context.publish({
          update: () => {
            liveModels = stored.models.filter((model) => model.provider === "openrouter");
          },
        });
      }
      if (!context.allowNetwork || context.signal.aborted) return;
      if (
        !context.force && stored?.etag === CACHE_MARKER && stored.checkedAt &&
        Date.now() - stored.checkedAt < REFRESH_INTERVAL_MS
      ) return;

      const response = await fetch(CATALOG_URL, {
        headers: { accept: "application/json" },
        signal: AbortSignal.any([context.signal, AbortSignal.timeout(10_000)]),
      });
      if (!response.ok) throw new Error(`OpenRouter catalog: HTTP ${response.status}`);
      const data: unknown = await response.json();
      if (typeof data !== "object" || data === null || !("data" in data) || !Array.isArray(data.data)) {
        throw new Error("OpenRouter catalog: invalid response");
      }
      const models = data.data.flatMap((entry: unknown) => {
        if (typeof entry !== "object" || entry === null || typeof (entry as OpenRouterModel).id !== "string" || typeof (entry as OpenRouterModel).name !== "string") return [];
        const model = catalogModel(entry as OpenRouterModel);
        return model ? [model] : [];
      });
      if (!models.length) throw new Error("OpenRouter catalog: no text models returned");
      await context.publish({
        persist: { models, checkedAt: Date.now(), etag: CACHE_MARKER },
        update: () => { liveModels = models; },
      });
    },
  });

  pi.on("session_start", async (_event, ctx) => {
    const result = await ctx.modelRegistry.refresh({ providers: ["openrouter"] });
    const error = result.errors.get("openrouter");
    if (error) ctx.ui.notify(`OpenRouter catalog refresh failed: ${error.message}`, "warning");
  });
}