import { getAgentDir, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
import { join } from "node:path";
import { streamSimple } from "./stream.js";
import { registerParleyFooter } from "./footer.js";
import { catalogScope, catalogSummary, discoverModels, type CatalogResult } from "./catalog.js";
import {
  PARLEY_API_KEY_ENV_VARS,
  pickStartupProfile,
  registerParleyKeySwitcher,
  resolveParleyApiKey,
  resolveParleyBaseUrl,
} from "./keys.js";

export function getParleyApiKeyEnvVar(): string {
  return pickStartupProfile()?.apiKeyVar ?? PARLEY_API_KEY_ENV_VARS[0];
}

export function getParleyBaseUrl(): string {
  return resolveParleyBaseUrl();
}

export default async function (pi: ExtensionAPI) {
  const cacheDir = join(getAgentDir(), "cache", "mit-parley");
  const scope = () => {
    try { return catalogScope(getParleyBaseUrl(), resolveParleyApiKey()); }
    catch { return "invalid-url"; }
  };
  let activeScope = scope();
  let generation = 0;
  let catalog: CatalogResult;

  const refresh = async (forceDocs = false, signal?: AbortSignal) => {
    const version = ++generation;
    const requestedScope = scope();
    const result = await discoverModels({
      baseUrl: getParleyBaseUrl(),
      apiKey: process.env.PARLEY_MODEL_DISCOVERY === "0" ? undefined : resolveParleyApiKey(),
      cacheDir,
      docsEnabled: process.env.PARLEY_DOCS_REFRESH !== "0",
      forceDocs,
      signal,
    });
    // A slow refresh for a previous key must not replace the active key's list.
    if (version === generation && requestedScope === scope()) {
      catalog = result;
      activeScope = requestedScope;
    }
    return catalog;
  };

  const install = () => pi.registerProvider("MIT-parley", {
    name: "MIT Parley",
    baseUrl: getParleyBaseUrl(),
    apiKey: `$${getParleyApiKeyEnvVar()}`,
    api: "parley-chat-completions",
    streamSimple,
    models: catalog.models,
    refreshModels: async context => (await refresh(true, context.signal)).models,
  });

  const refreshSelection = async (ctx: ExtensionContext) => {
    if (ctx.model?.provider !== "MIT-parley") return;
    const selected = ctx.modelRegistry.find("MIT-parley", ctx.model.id);
    if (selected) await pi.setModel(selected);
    else if (ctx.hasUI) ctx.ui.notify("The selected model is absent from this Parley catalog. Choose another with /model.", "warning");
  };

  // Factory initialization is awaited by Pi before startup model selection and
  // --list-models. Do not move discovery to session_start.
  await refresh();
  install();
  registerParleyKeySwitcher(pi, async ctx => {
    if (scope() === activeScope) return;
    await refresh();
    install();
    await refreshSelection(ctx);
  });
  registerParleyFooter(pi);

  pi.on("session_start", (_event, ctx) => {
    if (ctx.hasUI && ctx.model?.provider === "MIT-parley" && catalog.warnings.length) {
      ctx.ui.notify(catalog.warnings.join("\n") + " Use /parley-models for details.", "warning");
    }
  });
  pi.registerCommand("parley-models", {
    description: "Show Parley catalog sources, or refresh models, prices and limits: /parley-models refresh",
    handler: async (args, ctx) => {
      if (args.trim() && args.trim() !== "refresh") {
        ctx.ui.notify("Usage: /parley-models [refresh]", "info");
        return;
      }
      if (args.trim() === "refresh") {
        await ctx.waitForIdle();
        await refresh(true);
        install();
        await refreshSelection(ctx);
      }
      ctx.ui.notify(catalogSummary(catalog), catalog.warnings.length ? "warning" : "info");
    },
  });
}