import type { ExtensionAPI, ExtensionContext, ProviderModelConfig } from "@earendil-works/pi-coding-agent";
import type { OAuthCredentials, OAuthLoginCallbacks } from "@earendil-works/pi-ai";

const ROOT_URL = "https://vip.j3gb.com";
const OPENAI_URL = `${ROOT_URL}/v1`;
const REFRESH_INTERVAL_MS = 5 * 60_000;
const MODEL_DEFAULTS = {
  reasoning: false,
  input: ["text"] as ("text" | "image")[],
  cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
  contextWindow: 128_000,
  maxTokens: 8_192,
};

// j3gb and j3gb-alt keep their original Anthropic Messages behavior.
const SLOTS = [
  { id: "j3gb", name: "J3GB", env: "J3GB_API_KEY", api: "anthropic-messages", baseUrl: ROOT_URL },
  { id: "j3gb-alt", name: "J3GB (Alt)", env: "J3GB_ALT_API_KEY", api: "anthropic-messages", baseUrl: ROOT_URL },
  { id: "j3gb-anthropic", name: "J3GB (Anthropic)", env: "J3GB_ANTHROPIC_API_KEY", api: "anthropic-messages", baseUrl: ROOT_URL },
  { id: "j3gb-openai", name: "J3GB (OpenAI)", env: "J3GB_OPENAI_API_KEY", api: "openai-completions", baseUrl: OPENAI_URL },
] as const;

type Slot = (typeof SLOTS)[number];
type ProviderName = Slot["id"];

function slotFor(id: ProviderName): Slot {
  const slot = SLOTS.find((item) => item.id === id);
  if (!slot) throw new Error(`Unknown J3GB provider: ${id}`);
  return slot;
}

function isProviderName(value: string): value is ProviderName {
  return SLOTS.some((slot) => slot.id === value);
}

function model(id: string, name = id): ProviderModelConfig {
  return { id, name, ...MODEL_DEFAULTS };
}

function fallbackModels(slot: Slot): ProviderModelConfig[] {
  return slot.api === "anthropic-messages"
    ? [model("claude-haiku-4-5", "Claude Haiku 4.5")]
    : [model("15/gpt-5.4", "GPT-5.4")];
}

// Both wire APIs accept the same real model IDs. Grok is deliberately excluded.
function accepts(item: { id: string; display_name?: string }): boolean {
  return !`${item.id} ${item.display_name ?? ""}`.toLowerCase().includes("grok");
}

async function fetchModels(key: string): Promise<ProviderModelConfig[] | undefined> {
  try {
    const response = await fetch(`${OPENAI_URL}/models`, {
      headers: { Authorization: `Bearer ${key}` },
      signal: AbortSignal.timeout(10_000),
    });
    if (!response.ok) return;

    const payload = (await response.json()) as { data?: Array<{ id: string; display_name?: string }> };
    const models = (payload.data ?? [])
      .filter(accepts)
      .map((item) => model(item.id, item.display_name ?? item.id));
    return models.length ? models : undefined;
  } catch {
    return;
  }
}

function shortError(error?: string): string {
  const text = (error || "Request failed").replace(/\s+/g, " ").trim();
  return text.length > 160 ? `${text.slice(0, 157)}…` : text;
}

export default function (pi: ExtensionAPI) {
  let disposed = false;
  const inFlight = new Map<ProviderName, Promise<boolean>>();
  const lastRefresh = new Map<ProviderName, number>();
  const refreshing = new Set<ProviderName>();
  const errors = new Map<ProviderName, string>();

  function renderWidget(ctx: ExtensionContext) {
    if (!ctx.hasUI) return;
    if (refreshing.size) {
      const names = [...refreshing].map((id) => slotFor(id).name).join(", ");
      ctx.ui.setWidget("j3gb-provider", [`↻ J3GB: refreshing model catalog (${names})…`], { placement: "aboveEditor" });
    } else if (errors.size) {
      ctx.ui.setWidget(
        "j3gb-provider",
        ["⚠ J3GB provider error", ...[...errors].map(([id, error]) => `${slotFor(id).name}: ${error}`)],
        { placement: "aboveEditor" },
      );
    } else {
      ctx.ui.setWidget("j3gb-provider", undefined);
    }
  }

  function register(slot: Slot, models = fallbackModels(slot)) {
    pi.registerProvider(slot.id, {
      name: slot.name,
      baseUrl: slot.baseUrl,
      apiKey: `$${slot.env}`,
      api: slot.api,
      models,
      oauth: {
        name: slot.name,
        async login(callbacks: OAuthLoginCallbacks): Promise<OAuthCredentials> {
          const access = (await callbacks.onPrompt({ message: "Paste your J3GB API key (sk-...):" })).trim();
          if (!access.startsWith("sk-")) throw new Error("J3GB API keys must start with sk-.");
          await refresh(slot.id, access, true);
          return { access, refresh: access, expires: Date.now() + 365 * 24 * 60 * 60 * 1000 };
        },
        async refreshToken(credentials: OAuthCredentials) {
          // J3GB keys are opaque API keys; renew Pi's local credential lease.
          return { ...credentials, expires: Date.now() + 365 * 24 * 60 * 60 * 1000 };
        },
        getApiKey(credentials: OAuthCredentials) {
          return credentials.access;
        },
      },
    });
  }

  async function refresh(id: ProviderName, key: string, force = false): Promise<boolean> {
    if (disposed) return false;
    const existing = inFlight.get(id);
    if (existing) return existing;
    if (!force && Date.now() - (lastRefresh.get(id) ?? 0) < REFRESH_INTERVAL_MS) return false;

    const task = (async () => {
      const models = await fetchModels(key);
      if (disposed || !models) return false;
      register(slotFor(id), models);
      lastRefresh.set(id, Date.now());
      return true;
    })().catch(() => false).finally(() => inFlight.delete(id));

    inFlight.set(id, task);
    return task;
  }

  async function refreshSaved(id: ProviderName, ctx: ExtensionContext, force = false) {
    let key: string | undefined;
    try {
      key = await ctx.modelRegistry.getApiKeyForProvider(id);
    } catch {
      return;
    }
    if (!key || (!force && !inFlight.has(id) && Date.now() - (lastRefresh.get(id) ?? 0) < REFRESH_INTERVAL_MS)) return;

    refreshing.add(id);
    renderWidget(ctx);
    try {
      const updated = await refresh(id, key, force);
      if (force && !updated && !errors.has(id)) errors.set(id, "Could not refresh catalog; using fallback models.");
    } finally {
      refreshing.delete(id);
      renderWidget(ctx);
    }
  }

  SLOTS.forEach((slot) => register(slot));

  pi.on("session_start", async (_event, ctx) => {
    await Promise.all(SLOTS.map((slot) => refreshSaved(slot.id, ctx, true)));
  });

  pi.on("session_shutdown", (_event, ctx) => {
    disposed = true;
    inFlight.clear();
    refreshing.clear();
    errors.clear();
    if (ctx.hasUI) ctx.ui.setWidget("j3gb-provider", undefined);
  });

  pi.on("message_end", async (event, ctx) => {
    if (event.message.role !== "assistant" || !isProviderName(event.message.provider)) return;
    const id = event.message.provider;
    if (event.message.stopReason === "error") {
      errors.set(id, shortError(event.message.errorMessage));
      renderWidget(ctx);
      await refreshSaved(id, ctx);
    } else if (event.message.stopReason !== "aborted") {
      errors.delete(id);
      renderWidget(ctx);
    }
  });
}