import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import { createRunArtifacts } from "./src/artifacts.ts";
import { monitorRun } from "./src/monitor.ts";
import { listAvailableModelReferences } from "./src/models.ts";
import { createTmuxSingleRunner } from "./src/runner.ts";
import {
  buildCurrentWindowArgs,
  buildKillPaneArgs,
  buildSplitWindowArgs,
  buildWrapperShellCommand,
  isInsideTmux,
} from "./src/tmux.ts";
import { createSubagentParamsSchema } from "./src/schema.ts";
import { createSubagentTool } from "./src/tool.ts";

const packageRoot = dirname(fileURLToPath(import.meta.url));
const wrapperPath = join(packageRoot, "src", "wrapper", "cli.mjs");

export default function tmuxSubagentExtension(pi: ExtensionAPI) {
  if (process.env.PI_TMUX_SUBAGENT_GITHUB_COPILOT_INITIATOR === "agent") {
    pi.registerProvider("github-copilot", {
      headers: { "X-Initiator": "agent" },
    });
  }

  // In wrapper/child sessions spawned by the tmux runner we must not register the
  // subagent tool (that would cause nested subagent registrations). Skip all
  // subagent-tool registration logic when PI_TMUX_SUBAGENT_CHILD is set. Provider
  // overrides (above) are still allowed in child runs, so the guard is placed
  // after provider registration.
  if (process.env.PI_TMUX_SUBAGENT_CHILD === "1") {
    return;
  }

  let lastRegisteredModelsKey: string | undefined;

  const runSingleTask = createTmuxSingleRunner({
    assertInsideTmux() {
      if (!isInsideTmux()) throw new Error("tmux-backed subagents require pi to be running inside tmux.");
    },
    async getCurrentWindowId() {
      const result = await pi.exec("tmux", buildCurrentWindowArgs());
      return result.stdout.trim();
    },
    createArtifacts: createRunArtifacts,
    buildWrapperCommand(metaPath: string) {
      return buildWrapperShellCommand({ nodePath: process.execPath, wrapperPath, metaPath });
    },
    async createPane(input) {
      const result = await pi.exec("tmux", buildSplitWindowArgs(input));
      return result.stdout.trim();
    },
    monitorRun,
    async killPane(paneId: string) {
      await pi.exec("tmux", buildKillPaneArgs(paneId));
    },
  });

  const registerSubagentTool = (availableModels: string[]) => {
    // Do not register a tool when no models are available. Remember that the
    // last-registered key is different from the empty sentinel so that a later
    // non-empty list will still trigger registration.
    if (!availableModels || availableModels.length === 0) {
      const emptyKey = "\u0000";
      if (lastRegisteredModelsKey === emptyKey) return;
      lastRegisteredModelsKey = emptyKey;
      return;
    }

    // Create a deduplication key that is independent of the order of
    // availableModels by sorting a lowercase copy. Do not mutate
    // availableModels itself since we want to preserve the original order for
    // schema enum values.
    const key = [...availableModels].map((s) => s.toLowerCase()).sort().join("\u0000");
    if (key === lastRegisteredModelsKey) return;

    lastRegisteredModelsKey = key;
    pi.registerTool(
      createSubagentTool({
        parameters: createSubagentParamsSchema(availableModels),
        runSingleTask,
      }),
    );
  };

  const syncSubagentTool = (ctx: { modelRegistry: { getAvailable(): Array<{ provider: string; id: string }> } }) => {
    registerSubagentTool(listAvailableModelReferences(ctx.modelRegistry));
  };

  pi.on("session_start", (_event, ctx) => {
    syncSubagentTool(ctx);
  });

  pi.on("before_agent_start", (_event, ctx) => {
    syncSubagentTool(ctx);
  });
}