import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { DEFAULT_CONFIG, loadConfig } from "../src/config.js";
import { route } from "../src/router.js";
import { collectRepositorySignals } from "../src/signals.js";
import type { AvailableModel, ConfigResult, RepositorySignals, RouteDecision, RouteOverride } from "../src/types.js";

type RouterDependencies = {
  loadConfig?: (cwd: string) => ConfigResult;
  collectSignals?: (cwd: string) => RepositorySignals;
};

type RegistryModel = { provider: string; id: string; name?: string };

const modelKey = (model: RegistryModel) => `${model.provider}/${model.id}`;
const routeKey = (decision: RouteDecision) => `${decision.model ?? ""}/${decision.thinking}`;

function notify(ctx: { ui: { notify(message: string, type?: "info" | "warning" | "error"): void } }, message: string, type: "info" | "warning" = "info"): void {
  try {
    ctx.ui.notify(message, type);
  } catch {
    // Pi UI errors must never affect routing.
  }
}

export function createRouterExtension(pi: ExtensionAPI, dependencies: RouterDependencies = {}): void {
  const getConfig = dependencies.loadConfig ?? loadConfig;
  const getSignals = dependencies.collectSignals ?? collectRepositorySignals;
  let cachedConfig: { cwd: string; result: ConfigResult } | undefined;
  let cachedModels: { registry: unknown; expiresAt: number; models: RegistryModel[] } | undefined;
  let sessionEnabled: boolean | undefined;
  let override: RouteOverride | undefined;
  let manualSelection = false;
  let applyingRoute = false;
  let ignoreNextThinkingSelection = false;
  let lastDecision: RouteDecision | undefined;
  let lastRouteKey: string | undefined;
  const announcedWarnings = new Set<string>();

  const configFor = (cwd: string): ConfigResult => {
    if (cachedConfig?.cwd === cwd) return cachedConfig.result;
    try {
      const result = getConfig(cwd);
      cachedConfig = { cwd, result };
      return result;
    } catch {
      const result = { config: DEFAULT_CONFIG, warnings: ["configuration unavailable"] };
      cachedConfig = { cwd, result };
      return result;
    }
  };

  const availableModels = (ctx: { modelRegistry: { getAvailable(): RegistryModel[] } }): RegistryModel[] => {
    if (cachedModels?.registry === ctx.modelRegistry && cachedModels.expiresAt > Date.now()) return cachedModels.models;
    try {
      const models = ctx.modelRegistry.getAvailable();
      cachedModels = { registry: ctx.modelRegistry, expiresAt: Date.now() + 10_000, models };
      return models;
    } catch {
      return [];
    }
  };

  const announce = (ctx: any, decision: RouteDecision, debug: boolean): void => {
    const model = decision.model ? ` (${decision.model})` : "";
    const reason = debug ? `; ${decision.reasons[0]}` : "";
    notify(ctx, `Router: ${decision.profile}/${decision.thinking}${model}${reason}`);
  };

  pi.on("model_select", () => {
    if (!applyingRoute) manualSelection = true;
  });
  pi.on("thinking_level_select", () => {
    if (ignoreNextThinkingSelection) {
      ignoreNextThinkingSelection = false;
    } else if (!applyingRoute) {
      manualSelection = true;
    }
  });

  pi.on("before_agent_start", async (event, ctx) => {
    try {
      if (manualSelection) return;
      const loaded = configFor(ctx.cwd);
      for (const warning of loaded.warnings) {
        if (!announcedWarnings.has(warning)) {
          announcedWarnings.add(warning);
          notify(ctx, `Router: configuration warning: ${warning}`, "warning");
        }
      }
      const enabled = sessionEnabled ?? loaded.config.enabled;
      if (!enabled && !override) return;
      const models = availableModels(ctx);
      const overriddenModel = override?.model;
      if (overriddenModel && !models.some((model) => modelKey(model) === overriddenModel)) {
        override = undefined;
        notify(ctx, "Router: unavailable override cleared", "warning");
      }
      if (!enabled && !override) return;
      const decision = route({
        prompt: event.prompt,
        config: { ...loaded.config, enabled },
        availableModels: models.map((model): AvailableModel => ({ id: modelKey(model), name: model.name })),
        currentModel: ctx.model ? modelKey(ctx.model as RegistryModel) : undefined,
        repository: getSignals(ctx.cwd),
        session: lastDecision ? { current: lastDecision, taskBoundary: true } : undefined,
        override,
      });
      if (!decision.model) return;
      const selected = models.find((model) => modelKey(model) === decision.model);
      if (!selected) return;
      const key = routeKey(decision);
      const changed = key !== lastRouteKey;
      let appliedDecision = decision;

      if (changed) {
        applyingRoute = true;
        try {
          if (!ctx.model || modelKey(ctx.model as RegistryModel) !== decision.model) {
            const applied = await pi.setModel(selected as Parameters<ExtensionAPI["setModel"]>[0]);
            if (!applied) {
              notify(ctx, "Router: model change rejected; retained current selection", "warning");
              return;
            }
          }
          const currentThinking = (pi as { getThinkingLevel?: () => RouteDecision["thinking"] }).getThinkingLevel?.();
          if (currentThinking !== decision.thinking) {
            ignoreNextThinkingSelection = true;
            pi.setThinkingLevel(decision.thinking);
          }
          appliedDecision = {
            ...decision,
            thinking: (pi as { getThinkingLevel?: () => RouteDecision["thinking"] }).getThinkingLevel?.() ?? decision.thinking,
          };
        } finally {
          applyingRoute = false;
        }
      }
      lastDecision = appliedDecision;
      lastRouteKey = routeKey(appliedDecision);
      announce(ctx, appliedDecision, loaded.config.debug);
    } catch {
      // All adapter failures fail open: Pi retains its existing selection and turn.
    }
  });

  pi.registerCommand("codex-router", {
    description: "Show or control Codex routing: status, enable, disable, override, clear",
    handler: async (args, ctx) => {
      try {
        const [command = "status", value, thinking] = args.trim().split(/\s+/);
        if (command === "enable") {
          sessionEnabled = true;
          manualSelection = false;
          notify(ctx, "Router: enabled");
        } else if (command === "disable") {
          sessionEnabled = false;
          notify(ctx, "Router: disabled");
        } else if (command === "clear") {
          override = undefined;
          notify(ctx, "Router: override cleared");
        } else if (command === "override" && value) {
          const levels = new Set(["off", "minimal", "low", "medium", "high", "xhigh", "max"]);
          override = value === "fast" || value === "balanced" || value === "advanced" ? { profile: value } : { model: value };
          if (thinking && levels.has(thinking)) override.thinking = thinking as RouteOverride["thinking"];
          manualSelection = false;
          notify(ctx, `Router: override ${value}${thinking ? `/${thinking}` : ""}`);
        } else if (command === "status") {
          const mode = sessionEnabled === false ? "disabled" : manualSelection ? "manual" : "enabled";
          const current = lastDecision
            ? `${lastDecision.profile}/${lastDecision.thinking}${lastDecision.model ? ` (${lastDecision.model})` : ""}`
            : "pending";
          notify(ctx, `Router: ${mode}; ${current}${override ? "; override" : ""}`);
        } else {
          notify(ctx, "Router: use status, enable, disable, override <profile|model> [thinking], or clear");
        }
      } catch {
        // Commands must remain non-fatal even if the host UI fails.
      }
    },
  });
}

export default function routerExtension(pi: ExtensionAPI): void {
  createRouterExtension(pi);
}