import { StringEnum } from "@earendil-works/pi-ai";
import {
  getAgentDir,
  type ExtensionAPI,
  type ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";

import {
  loadAgentShellConfig,
  setRosterEnabled,
  type AgentShellInteractiveConfig,
  type AgentShellRosterConfig,
} from "./config.ts";
import {
  JobRegistry,
  type JobActivity,
  type JobSnapshot,
  type JobStatus,
} from "./jobs.ts";
import type { AgentShellLimits } from "./limits.ts";
import {
  AGENT_SHELL_PROJECT_DIRECTORY,
  closeAgentShellWorkers,
  getAgentShellModels,
  getSupportedAgentTypes,
  isAgentShellRuntimeInstalled,
  runAgentShell,
  supportsAgentShellInteractive,
  supportsAgentShellModelDiscovery,
  type RunResult,
} from "./runner.ts";

const UV_INSTALL_URL =
  "https://docs.astral.sh/uv/getting-started/installation/";
const OUTPUT_MODE_ENTRY_TYPE = "agentshell-output-mode";
const JOB_RESULT_MESSAGE_TYPE = "agentshell-job-result";
const JOB_WIDGET_KEY = "agentshell-jobs";
const JOB_WIDGET_OPTIONS = { placement: "aboveEditor" as const };
const MAX_WIDGET_ACTIVITY_CHARACTERS = 40;
const MAX_INSPECTION_ACTIVITY_ENTRIES = 20;
const MAX_INSPECTION_ENTRY_CHARACTERS = 400;
const RESUME_SESSION_GUIDANCE = [
  "Omit `resume_session_id` for a new session.",
  "If the tool-call interface requires this property, use `null` for a new session.",
  "For a resumed session, the value must come from an earlier successful " +
    "subagent result.",
  "Do not pass `new`, a background Job ID, or a newly generated UUID.",
].join(" ");

type TerminalJobStatus = Exclude<JobStatus, "running">;

interface JobResultMessageDetails {
  jobId: string;
  status: TerminalJobStatus;
  warnings: string[];
  silent?: boolean;
}

interface PendingTerminalMessage {
  jobId: string;
  status: TerminalJobStatus;
  output: string;
  warnings: string[];
  silent: boolean;
}

function setupCommand(): string {
  return [
    "uv sync --project",
    `"${AGENT_SHELL_PROJECT_DIRECTORY}"`,
    "--locked",
  ].join(" ");
}

interface InteractiveLaunchSettings {
  interactive: boolean;
  stay_open: boolean;
  inactivity_timeout: number;
  inactivity_enabled: boolean;
}

function resolveInteractiveLaunchSettings(
  params: Partial<InteractiveLaunchSettings>,
  config: AgentShellInteractiveConfig,
): InteractiveLaunchSettings {
  const interactive = params.interactive ?? config.enabled;

  return {
    interactive,
    stay_open: interactive
      ? params.stay_open ?? config.stay_open
      : false,
    inactivity_timeout: params.inactivity_timeout ??
      config.inactivity_timeout,
    inactivity_enabled: params.inactivity_enabled ??
      config.inactivity_enabled,
  };
}

function formatRunOutput(result: RunResult): string {
  const warningOutput = result.details.warnings
    .map((warning) => `Warning: ${warning}`)
    .join("\n");
  const sessionId = result.details.sessionId;
  const sessionOutput = sessionId ? `Session ID: ${sessionId}` : "";

  return [warningOutput, result.output, sessionOutput]
    .filter((part) => part.length > 0)
    .join("\n\n");
}

function formatJobLaunch(jobId: string) {
  return {
    content: [
      {
        type: "text" as const,
        text: `Subagent job ${jobId} started.`,
      },
    ],
    details: {
      status: "running" as const,
      jobId,
      outputTokens: 0,
      warnings: [] as string[],
    },
  };
}

function errorMessage(error: unknown): string {
  return error instanceof Error ? error.message : String(error);
}

function formatTerminalMessage(
  jobId: string,
  status: TerminalJobStatus,
  output: string,
): string {
  if (status === "completed") {
    return `Subagent job ${jobId} completed.\n\n${output}`;
  }

  if (status === "cancelled") {
    return `Subagent job ${jobId} was cancelled.\n\n${output}`;
  }

  return `Subagent job ${jobId} failed.\n\n${output}`;
}

function formatJobWidgetRow(
  job: JobSnapshot,
  isLast: boolean,
  deliveryStatus?: TerminalJobStatus,
): string {
  const branch = isLast ? "└─" : "├─";
  const taskName = job.taskName ?? "Subagent";
  const runtime = [
    job.agentType ?? "subagent",
    job.model ?? "default",
    job.effort ?? "default",
  ].join("/");
  const status = deliveryStatus === undefined
    ? job.status === "cancelled" ? " · cancelling…" : ""
    : ` · ${deliveryStatus} · delivering…`;

  return `${branch} ${taskName} · ${runtime} · ${job.id.slice(0, 12)}${status}`;
}

function cleanActivityText(value: string): string {
  return value
    .replace(/\u001b\[[0-?]*[ -/]*[@-~]/g, "")
    .replace(/[\u0000-\u0008\u000b\u000c\u000e-\u001f\u007f]/g, "")
    .replace(/\s+/g, " ")
    .trim();
}

function truncateActivityText(value: string, maxCharacters: number): string {
  const characters = Array.from(value);

  if (characters.length <= maxCharacters) {
    return value;
  }

  return `${characters.slice(0, maxCharacters - 1).join("")}…`;
}

function describeToolActivity(job: JobSnapshot, content: string): string {
  const cleaned = cleanActivityText(content);

  if (cleaned.length === 0) {
    return "using a tool";
  }

  const normalized = cleaned.toLowerCase().replace(/[^a-z0-9]/g, "");
  const knownDescriptions: Record<string, string> = {
    read: "reading a file",
    view: "reading a file",
    edit: "editing a file",
    write: "editing a file",
    notebookedit: "editing a notebook",
    grep: "searching the codebase",
    glob: "searching the codebase",
    websearch: "searching the web",
    webfetch: "fetching a web page",
    bash: "shell command",
    shell: "shell command",
  };
  const known = knownDescriptions[normalized];

  if (known !== undefined) {
    return known;
  }

  const shortContent = truncateActivityText(
    cleaned,
    MAX_WIDGET_ACTIVITY_CHARACTERS,
  );

  if (
    job.agentType === "codex" ||
    job.agentType === "cursor" ||
    /[\s|&;<>]/.test(cleaned)
  ) {
    return `shell command: \`${shortContent}\``;
  }

  return `using ${shortContent}`;
}

function describeLatestActivity(job: JobSnapshot): string | undefined {
  const activity = job.latestActivity;

  if (activity === undefined) {
    return undefined;
  }

  if (activity.type === "tool_use") {
    return describeToolActivity(job, activity.content);
  }

  if (activity.type === "text") {
    return "reporting progress";
  }

  return activity.type === "warning"
    ? "warning reported"
    : "error reported";
}

function formatActivityEntry(activity: JobActivity): string {
  const label = activity.type === "tool_use" ? "tool" : activity.type;
  const cleaned = cleanActivityText(activity.content);
  const content = truncateActivityText(
    cleaned,
    MAX_INSPECTION_ENTRY_CHARACTERS,
  );

  return `[${label}] ${content || "(no displayable content)"}`;
}

function formatJobInspection(
  job: JobSnapshot,
  isDelivering = false,
): string {
  const runtime = [
    job.agentType ?? "subagent",
    job.model ?? "default",
    job.effort ?? "default",
  ].join("/");

  if (isDelivering) {
    return [
      `Subagent job ${job.id}: delivering`,
      `Task: ${job.taskName ?? "Subagent"}`,
      `Agent: ${runtime}`,
      "Final result is waiting for delivery.",
    ].join("\n");
  }

  const latest = describeLatestActivity(job) ?? "waiting for activity";
  const activity = job.activity
    ?.slice(-MAX_INSPECTION_ACTIVITY_ENTRIES)
    .map(formatActivityEntry) ?? [];

  return [
    `Subagent job ${job.id}: ${job.status}`,
    `Task: ${job.taskName ?? "Subagent"}`,
    `Agent: ${runtime}`,
    `Last activity: ${latest}`,
    "",
    "Activity:",
    ...(activity.length > 0 ? activity : ["No activity reported yet."]),
  ].join("\n");
}

function formatJobListEntry(
  job: JobSnapshot,
  isDelivering: boolean,
): string {
  const status = isDelivering ? "delivering" : job.status;
  const details = [
    job.taskName ?? "Subagent",
    job.agentType ?? "subagent",
    job.model ?? "default",
    `effort: ${job.effort ?? "default"}`,
  ].join(" · ");
  const latest = isDelivering ? undefined : describeLatestActivity(job);

  return [
    `${job.id}: ${status}`,
    `  ${details}`,
    ...(latest === undefined ? [] : [`  Last activity: ${latest}`]),
  ].join("\n");
}

function formatJobLookupFailure(
  jobId: string,
  jobs: JobRegistry,
): string {
  const activeJobIds = jobs.list().map((job) => `- ${job.id}`);
  const activeJobs = activeJobIds.length === 0
    ? "No subagent jobs are currently active."
    : ["Active subagent jobs:", ...activeJobIds].join("\n");

  return [
    `No active subagent job found with ID ${jobId}.`,
    activeJobs,
  ].join("\n\n");
}

function formatCancellationFailure(
  jobId: string,
  jobs: JobRegistry,
): string {
  const runningJobIds = jobs.list()
    .filter((job) => job.status === "running")
    .map((job) => `- ${job.id}`);
  const runningJobs = runningJobIds.length === 0
    ? "No subagent jobs are currently running."
    : ["Running subagent jobs:", ...runningJobIds].join("\n");

  return [
    `No running subagent job found with ID ${jobId}.`,
    runningJobs,
  ].join("\n\n");
}

function updateJobWidget(
  ctx: Pick<ExtensionContext, "hasUI" | "ui">,
  jobs: JobRegistry,
  deliveries: ReadonlyMap<string, TerminalJobStatus>,
): void {
  if (!ctx.hasUI) {
    return;
  }

  const activeJobs = jobs.list();
  if (activeJobs.length === 0) {
    ctx.ui.setWidget(JOB_WIDGET_KEY, undefined, JOB_WIDGET_OPTIONS);
    return;
  }

  const running = activeJobs.filter((job) => job.status === "running").length;
  const cancelling = activeJobs.filter((job) =>
    job.status === "cancelled" && !deliveries.has(job.id)
  ).length;
  const delivering = activeJobs.filter((job) => deliveries.has(job.id)).length;
  const counts = [
    running > 0 ? `${running} running` : "",
    cancelling > 0 ? `${cancelling} cancelling` : "",
    delivering > 0 ? `${delivering} delivering` : "",
  ].filter((part) => part.length > 0);
  const glyph = running > 0 ? "●" : cancelling > 0 ? "■" : "◆";
  const header = ctx.ui.theme.fg(
    cancelling > 0 && running === 0 ? "warning" : "accent",
    `${glyph} Background agents · ${counts.join(" · ")}`,
  );
  const visibleJobs = activeJobs.slice(0, 3);
  const hiddenJobs = activeJobs.slice(visibleJobs.length);
  const rows: string[] = [];

  visibleJobs.forEach((job, index) => {
    const isLast = hiddenJobs.length === 0 && index === visibleJobs.length - 1;
    const deliveryStatus = deliveries.get(job.id);
    const row = formatJobWidgetRow(job, isLast, deliveryStatus);
    const styledRow = deliveryStatus === "completed"
      ? ctx.ui.theme.fg("success", row)
      : deliveryStatus === "failed"
      ? ctx.ui.theme.fg("error", row)
      : job.status === "cancelled"
      ? ctx.ui.theme.fg("warning", row)
      : row;

    rows.push(styledRow);

    const latest = describeLatestActivity(job);
    if (latest !== undefined && deliveryStatus === undefined) {
      const continuation = isLast ? "  " : "│ ";
      rows.push(
        ctx.ui.theme.fg(
          "dim",
          `${continuation} Last activity: ${latest}`,
        ),
      );
    }
  });

  if (hiddenJobs.length > 0) {
    const hiddenStatus = hiddenJobs.every((job) => job.status === "running")
      ? "running"
      : hiddenJobs.every((job) => deliveries.has(job.id))
      ? "delivering"
      : "active";
    rows.push(
      ctx.ui.theme.fg(
        "dim",
        `└─ +${hiddenJobs.length} more ${hiddenStatus} · /agentshell-jobs`,
      ),
    );
  }

  ctx.ui.setWidget(
    JOB_WIDGET_KEY,
    [header, ...rows],
    JOB_WIDGET_OPTIONS,
  );
}

function safelyUpdateJobWidget(
  ctx: Pick<ExtensionContext, "hasUI" | "ui">,
  jobs: JobRegistry,
  deliveries: ReadonlyMap<string, TerminalJobStatus>,
): void {
  try {
    updateJobWidget(ctx, jobs, deliveries);
  } catch (error) {
    process.stderr.write(
      `Could not update AgentShell job widget: ${errorMessage(error)}\n`,
    );
  }
}

function deliverTerminalMessage(
  pi: ExtensionAPI,
  jobs: JobRegistry,
  message: PendingTerminalMessage,
  isShuttingDown: () => boolean,
  deliveries: Map<string, TerminalJobStatus>,
  onWidgetChange: () => void,
): void {
  const { jobId, status, output, warnings, silent } = message;

  if (isShuttingDown() || jobs.get(jobId) === undefined) {
    deliveries.delete(jobId);
    jobs.remove(jobId);
    return;
  }

  try {
    pi.sendMessage<JobResultMessageDetails>(
      {
        customType: JOB_RESULT_MESSAGE_TYPE,
        content: formatTerminalMessage(jobId, status, output),
        display: true,
        details: {
          jobId,
          status,
          warnings,
          ...(silent ? { silent: true } : {}),
        },
      },
      {
        triggerTurn: true,
        deliverAs: "followUp",
      },
    );
    onWidgetChange();
  } catch (error) {
    process.stderr.write(
      `Could not deliver subagent job ${jobId}: ${errorMessage(error)}\n`,
    );
    deliveries.delete(jobId);
    jobs.remove(jobId);
    onWidgetChange();
  }
}

function sendTerminalMessage(
  pi: ExtensionAPI,
  jobs: JobRegistry,
  message: PendingTerminalMessage,
  isShuttingDown: () => boolean,
  isParentAgentActive: () => boolean,
  deliveries: Map<string, TerminalJobStatus>,
  pendingMessages: PendingTerminalMessage[],
  onWidgetChange: () => void,
): void {
  if (isShuttingDown() || jobs.get(message.jobId) === undefined) {
    jobs.remove(message.jobId);
    return;
  }

  deliveries.set(message.jobId, message.status);

  if (isParentAgentActive()) {
    pendingMessages.push(message);
    onWidgetChange();
    return;
  }

  deliverTerminalMessage(
    pi,
    jobs,
    message,
    isShuttingDown,
    deliveries,
    onWidgetChange,
  );
}

function flushPendingTerminalMessages(
  pi: ExtensionAPI,
  jobs: JobRegistry,
  pendingMessages: PendingTerminalMessage[],
  isShuttingDown: () => boolean,
  deliveries: Map<string, TerminalJobStatus>,
  onWidgetChange: () => void,
): void {
  const messages = pendingMessages.splice(0);

  for (const message of messages) {
    deliverTerminalMessage(
      pi,
      jobs,
      message,
      isShuttingDown,
      deliveries,
      onWidgetChange,
    );
  }
}

function registerJobDeliveryHandler(
  pi: ExtensionAPI,
  jobs: JobRegistry,
  deliveries: Map<string, TerminalJobStatus>,
): void {
  pi.on("message_start", (event, ctx) => {
    const message = event.message;

    if (
      message.role !== "custom" ||
      message.customType !== JOB_RESULT_MESSAGE_TYPE ||
      typeof message.details !== "object" ||
      message.details === null ||
      !("jobId" in message.details) ||
      typeof message.details.jobId !== "string"
    ) {
      return;
    }

    const jobId = message.details.jobId;

    if (deliveries.delete(jobId)) {
      jobs.remove(jobId);
      safelyUpdateJobWidget(ctx, jobs, deliveries);
    }
  });
}

function registerJobMessageRenderer(pi: ExtensionAPI): void {
  pi.registerMessageRenderer<JobResultMessageDetails>(
    JOB_RESULT_MESSAGE_TYPE,
    (message, _options, theme) => {
      const details = message.details;
      const content = typeof message.content === "string"
        ? message.content
        : message.content
          .filter((part) => part.type === "text")
          .map((part) => part.text)
          .join("\n");

      if (details?.silent && details.status === "completed") {
        const warnings = details.warnings.map((warning) =>
          theme.fg("warning", `Warning: ${warning}`)
        );
        const completed = theme.fg("success", "✓ Completed");
        return new Text([...warnings, completed].join("\n"), 0, 0);
      }

      const color = details?.status === "failed"
        ? "error"
        : details?.status === "cancelled"
        ? "warning"
        : "toolOutput";

      return new Text(theme.fg(color, content), 0, 0);
    },
  );
}

function registerSubagentModelsTool(
  pi: ExtensionAPI,
  agentTypes: string[],
  limits: AgentShellLimits,
): void {
  pi.registerTool({
    name: "subagent_list_models",
    label: "List subagent models",
    description: [
      "List the exact model selectors advertised by an AgentShell agent.",
      "Pass a returned selector unchanged to the subagent tool's model parameter.",
      "An empty list is valid and does not prove a model can run.",
    ].join(" "),
    parameters: Type.Object({
      agent_type: StringEnum(agentTypes, {
        description: "AgentShell agent type to inspect",
      }),
      cwd: Type.Optional(Type.String({
        minLength: 1,
        description:
          "Working directory used for workspace-aware discovery. " +
          "Defaults to Pi's current working directory.",
      })),
    }),

    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
      const models = await getAgentShellModels(
        params.agent_type,
        params.cwd ?? ctx.cwd,
        signal,
        limits,
      );

      return {
        content: [{
          type: "text" as const,
          text: JSON.stringify(models),
        }],
        details: {
          status: "ok" as const,
          agentType: params.agent_type,
          models,
          outputTokens: 0,
          warnings: [] as string[],
        },
      };
    },
  });
}

function registerSubagentStatusTool(
  pi: ExtensionAPI,
  jobs: JobRegistry,
  deliveries: ReadonlyMap<string, TerminalJobStatus>,
): void {
  pi.registerTool({
    name: "subagent_status",
    label: "Inspect subagent",
    description: [
      "Inspect the current status and recent activity of an active subagent job.",
      "Use this when progress is useful; do not repeatedly poll the job.",
      "Activity is held only in memory and is removed after final delivery.",
    ].join(" "),
    parameters: Type.Object({
      job_id: Type.String({
        minLength: 1,
        description: "Job ID returned by the subagent tool",
      }),
    }),

    async execute(_toolCallId, params) {
      const job = jobs.get(params.job_id);

      if (job === undefined) {
        throw new Error(formatJobLookupFailure(params.job_id, jobs));
      }

      const isDelivering = deliveries.has(job.id);

      return {
        content: [{
          type: "text" as const,
          text: formatJobInspection(job, isDelivering),
        }],
        details: {
          status: isDelivering ? "delivering" as const : job.status,
          jobId: job.id,
          outputTokens: 0,
          warnings: [] as string[],
        },
      };
    },
  });
}

async function registerSubagentTool(
  pi: ExtensionAPI,
  limits: AgentShellLimits,
  interactiveConfig: AgentShellInteractiveConfig,
  jobs: JobRegistry,
  deliveries: Map<string, TerminalJobStatus>,
  pendingMessages: PendingTerminalMessage[],
  isParentAgentActive: () => boolean,
  isSilentMode: () => boolean,
  isShuttingDown: () => boolean,
  modelDiscoverySupported: boolean,
  interactiveSupported: boolean,
  rosterEnabled: boolean,
): Promise<(enabled: boolean) => void> {
  const agentTypes = await getSupportedAgentTypes(limits);

  registerJobMessageRenderer(pi);
  registerJobDeliveryHandler(pi, jobs, deliveries);
  registerSubagentStatusTool(pi, jobs, deliveries);

  const registerInvocation = (enabled: boolean): void => pi.registerTool({
    name: "subagent",
    label: "Subagent",
    description: [
      "Delegate a task to an AI coding agent in a separate context.",
      "Long-running subagent calls may return `Subagent job <job-id> started`.",
      "Do not repeatedly poll the process or run sleep commands.",
      "Use subagent_status only when current progress is useful.",
      "A Job ID is not a session ID.",
      "The extension automatically delivers the result using a follow-up turn.",
      "Continue other work or remain idle until notified.",
      "The real resumable session ID arrives with the completion result.",
      RESUME_SESSION_GUIDANCE,
      ...(enabled ? [
        "Before delegating, check subagent_roster for preferred roles and settings.",
        "Roster entries are advisory; pass the chosen agent_type, model, and effort explicitly.",
      ] : []),
    ].join(" "),
    parameters: Type.Object({
      agent_type: StringEnum(agentTypes, {
        description: "AgentShell agent type to run",
      }),
      task_name: Type.String({
        minLength: 1,
        maxLength: 40,
        description:
          "Short human-readable name for this delegated task, " +
          "such as Code review or QA check",
      }),
      cwd: Type.Optional(Type.String({
        minLength: 1,
        description:
          "Working directory for the subagent. " +
          "Defaults to Pi's current working directory.",
      })),
      prompt: Type.String({
        description: "Task for the subagent",
      }),
      model: Type.Optional(Type.String({
        minLength: 1,
        description: "Model identifier passed to AgentShell",
      })),
      effort: Type.Optional(Type.String({
        minLength: 1,
        description: "Reasoning effort passed to AgentShell",
      })),
      resume_session_id: Type.Optional(Type.Union(
        [
          Type.String({ minLength: 1 }),
          Type.Null(),
        ],
        { description: RESUME_SESSION_GUIDANCE },
      )),
      auto_approve: Type.Optional(Type.Boolean({
        description:
          "Allow the subagent to approve tool use automatically. " +
          "Defaults to false for headless runs. Interactive runs use " +
          "the agent's native permission prompts.",
      })),
      allowed_tools: Type.Optional(Type.Array(
        Type.String({ minLength: 1 }),
        {
          minItems: 1,
          description:
            "Tool names the subagent may use for headless runs, when " +
            "supported. Interactive runs use native permission controls " +
            "and reject unsupported options.",
        },
      )),
      disallowed_tools: Type.Optional(Type.Array(
        Type.String({ minLength: 1 }),
        {
          minItems: 1,
          description:
            "Tool names the subagent must not use for headless runs, when " +
            "supported. Interactive runs use native permission controls " +
            "and reject unsupported options.",
        },
      )),
      interactive: Type.Optional(Type.Boolean({
        description:
          "Run the agent's native interactive UI in a neighbouring tmux " +
          "pane. Requires TMUX and TMUX_PANE. Defaults to the configured " +
          "interactive.enabled value.",
      })),
      stay_open: Type.Optional(Type.Boolean({
        description:
          "Keep an interactive pane open after its result is delivered. " +
          "Applies only to interactive runs and defaults to the configured " +
          "interactive.stay_open value.",
      })),
      inactivity_timeout: Type.Optional(Type.Number({
        exclusiveMinimum: 0,
        description:
          "Positive number of seconds used to infer completion after an " +
          "inactive interactive pane when inactivity is enabled. Defaults " +
          "to the configured interactive.inactivity_timeout value.",
      })),
      inactivity_enabled: Type.Optional(Type.Boolean({
        description:
          "Enable automatic closure after inactivity_timeout for interactive " +
          "runs. Defaults to the configured interactive.inactivity_enabled " +
          "value.",
      })),
    }),

    renderCall(args, theme) {
      let text =
        theme.fg("toolTitle", theme.bold("Subagent ")) +
        theme.fg("accent", args.agent_type);

      if (args.model) {
        text += theme.fg("muted", ` · model: ${args.model}`);
      }

      if (args.effort) {
        text += theme.fg("muted", ` · effort: ${args.effort}`);
      }

      return new Text(text, 0, 0);
    },

    renderResult(result, _options, theme, context) {
      const output = result.content
        .filter((part) => part.type === "text")
        .map((part) => part.text)
        .join("\n");
      const color = context.isError ? "error" : "toolOutput";

      return new Text(theme.fg(color, output), 0, 0);
    },

    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      const silent = isSilentMode();
      const interactive = resolveInteractiveLaunchSettings(
        params,
        interactiveConfig,
      );

      if (interactive.interactive && !interactiveSupported) {
        throw new Error([
          "Interactive AgentShell runs need an updated runtime.",
          `Run: ${setupCommand()}`,
          "Then restart Pi or run /reload.",
        ].join("\n"));
      }

      const job = jobs.start(
        (signal, jobId) =>
          runAgentShell(
            {
              agent_type: params.agent_type,
              cwd: params.cwd ?? ctx.cwd,
              prompt: params.prompt,
              model: params.model,
              effort: params.effort,
              session_id: params.resume_session_id ?? undefined,
              auto_approve: params.auto_approve,
              allowed_tools: params.allowed_tools,
              disallowed_tools: params.disallowed_tools,
              interactive: interactive.interactive,
              stay_open: interactive.stay_open,
              inactivity_timeout: interactive.inactivity_timeout,
              inactivity_enabled: interactive.inactivity_enabled,
            },
            signal,
            (update) => {
              jobs.recordActivity(jobId, update.activity);

              if (!isShuttingDown()) {
                safelyUpdateJobWidget(ctx, jobs, deliveries);
              }
            },
            limits,
          ),
        {
          taskName: params.task_name,
          agentType: params.agent_type,
          model: params.model,
          effort: params.effort,
        },
      );

      safelyUpdateJobWidget(ctx, jobs, deliveries);

      void job.completion.then(
        (result) => {
          const status = jobs.get(job.id)?.status === "cancelled"
            ? "cancelled"
            : "completed";
          const output = status === "cancelled"
            ? "The worker stopped after cancellation."
            : formatRunOutput(result);

          sendTerminalMessage(
            pi,
            jobs,
            {
              jobId: job.id,
              status,
              output,
              warnings: result.details.warnings,
              silent,
            },
            isShuttingDown,
            isParentAgentActive,
            deliveries,
            pendingMessages,
            () => safelyUpdateJobWidget(ctx, jobs, deliveries),
          );
        },
        (error: unknown) => {
          const status = jobs.get(job.id)?.status === "cancelled"
            ? "cancelled"
            : "failed";

          sendTerminalMessage(
            pi,
            jobs,
            {
              jobId: job.id,
              status,
              output: errorMessage(error),
              warnings: [],
              silent,
            },
            isShuttingDown,
            isParentAgentActive,
            deliveries,
            pendingMessages,
            () => safelyUpdateJobWidget(ctx, jobs, deliveries),
          );
        },
      );

      return formatJobLaunch(job.id);
    },
  });

  registerInvocation(rosterEnabled);

  if (modelDiscoverySupported) {
    registerSubagentModelsTool(pi, agentTypes, limits);
  }

  pi.registerTool({
    name: "subagent_cancel",
    label: "Cancel subagent",
    description: "Cancel a running AgentShell subagent job by ID",
    parameters: Type.Object({
      job_id: Type.String({
        minLength: 1,
        description: "Job ID returned by the subagent tool",
      }),
    }),

    async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
      if (!jobs.cancel(params.job_id)) {
        throw new Error(formatCancellationFailure(params.job_id, jobs));
      }

      safelyUpdateJobWidget(ctx, jobs, deliveries);

      return {
        content: [
          {
            type: "text" as const,
            text: `Subagent job ${params.job_id} cancelled.`,
          },
        ],
        details: {
          status: "cancelled" as const,
          jobId: params.job_id,
          outputTokens: 0,
          warnings: [] as string[],
        },
      };
    },
  });

  return registerInvocation;
}

function registerRosterTool(pi: ExtensionAPI, roster: AgentShellRosterConfig): void {
  pi.registerTool({
    name: "subagent_roster",
    label: "Subagent roster",
    description: "List preferred subagent roles and their suggested AgentShell settings.",
    parameters: Type.Object({}),
    async execute() {
      return {
        content: [{
          type: "text" as const,
          text: roster.roles.length === 0
            ? "No preferred subagent roles are configured."
            : JSON.stringify(roster.roles, null, 2),
        }],
      };
    },
  });
}

export default async function subagentsExtension(
  pi: ExtensionAPI,
): Promise<void> {
  if (process.env.PI_AGENT_SHELL_CHILD === "1") {
    return;
  }

  const agentDirectory = getAgentDir();
  const config = loadAgentShellConfig(agentDirectory);
  const limits: AgentShellLimits = config;
  const jobs = new JobRegistry(limits.maxOutputBytes);
  const deliveries = new Map<string, TerminalJobStatus>();
  const pendingMessages: PendingTerminalMessage[] = [];
  let parentAgentActive = false;
  let shuttingDown = false;
  let silentMode = false;
  let refreshSubagentTool: ((enabled: boolean) => void) | undefined;

  if (config.roster.enabled) {
    registerRosterTool(pi, config.roster);
  }

  pi.on("agent_start", () => {
    parentAgentActive = true;
  });

  pi.on("agent_settled", (_event, ctx) => {
    if (!ctx.isIdle()) {
      parentAgentActive = true;
      return;
    }

    parentAgentActive = false;

    flushPendingTerminalMessages(
      pi,
      jobs,
      pendingMessages,
      () => shuttingDown,
      deliveries,
      () => safelyUpdateJobWidget(ctx, jobs, deliveries),
    );
  });

  pi.on("session_start", (_event, ctx) => {
    parentAgentActive = false;
    pendingMessages.length = 0;
    deliveries.clear();
    shuttingDown = false;
    silentMode = false;

    for (const entry of ctx.sessionManager.getBranch()) {
      if (
        entry.type === "custom" &&
        entry.customType === OUTPUT_MODE_ENTRY_TYPE &&
        typeof entry.data === "object" &&
        entry.data !== null &&
        "silent" in entry.data &&
        typeof entry.data.silent === "boolean"
      ) {
        silentMode = entry.data.silent;
      }
    }
  });

  pi.on("session_shutdown", async (_event, ctx) => {
    parentAgentActive = false;
    pendingMessages.length = 0;
    shuttingDown = true;
    jobs.cancelAll();
    await closeAgentShellWorkers();
    jobs.clear();
    deliveries.clear();
    safelyUpdateJobWidget(ctx, jobs, deliveries);
  });

  pi.registerCommand("agentshell-silent", {
    description: "Toggle display of subagent responses",
    handler: async (_args, ctx) => {
      silentMode = !silentMode;
      pi.appendEntry(OUTPUT_MODE_ENTRY_TYPE, { silent: silentMode });
      ctx.ui.notify(
        silentMode
          ? "Subagent responses are now hidden."
          : "Subagent responses are now visible.",
        "info",
      );
    },
  });

  pi.registerCommand("agentshell-roster", {
    description: "Toggle the global preferred subagent roster",
    handler: async (_args, ctx) => {
      try {
        const currentRoster = loadAgentShellConfig(agentDirectory).roster;
        const enabled = !currentRoster.enabled;
        setRosterEnabled(agentDirectory, enabled);
        config.roster = { ...currentRoster, enabled };

        if (enabled) {
          registerRosterTool(pi, config.roster);
        }

        refreshSubagentTool?.(enabled);
        const activeTools = pi.getActiveTools();
        pi.setActiveTools(enabled
          ? [...new Set([...activeTools, "subagent_roster"])]
          : activeTools.filter((name) => name !== "subagent_roster"));
        ctx.ui.notify(
          enabled ? "Subagent roster is now enabled."
            : "Subagent roster is now disabled.",
          "info",
        );
      } catch (error) {
        ctx.ui.notify(`Could not toggle subagent roster: ${errorMessage(error)}`, "error");
      }
    },
  });

  pi.registerCommand("agentshell-jobs", {
    description: "List active AgentShell subagent jobs",
    handler: async (_args, ctx) => {
      const activeJobs = jobs.list();

      if (activeJobs.length === 0) {
        ctx.ui.notify("No active subagent jobs.", "info");
        return;
      }

      ctx.ui.notify(
        activeJobs
          .map((job) => formatJobListEntry(job, deliveries.has(job.id)))
          .join("\n"),
        "info",
      );
    },
  });

  pi.registerCommand("agentshell-inspect", {
    description: "Inspect recent activity for an active subagent job",
    handler: async (args, ctx) => {
      const jobId = args.trim();

      if (jobId.length === 0) {
        ctx.ui.notify("Usage: /agentshell-inspect <job-id>", "warning");
        return;
      }

      const job = jobs.get(jobId);
      ctx.ui.notify(
        job === undefined
          ? formatJobLookupFailure(jobId, jobs)
          : formatJobInspection(job, deliveries.has(jobId)),
        job === undefined ? "warning" : "info",
      );
    },
  });

  pi.registerCommand("agentshell-cancel", {
    description: "Cancel an active AgentShell subagent job by ID",
    handler: async (args, ctx) => {
      const jobId = args.trim();

      if (jobId.length === 0) {
        ctx.ui.notify("Usage: /agentshell-cancel <job-id>", "warning");
        return;
      }

      const cancelled = jobs.cancel(jobId);

      if (cancelled) {
        safelyUpdateJobWidget(ctx, jobs, deliveries);
      }

      ctx.ui.notify(
        cancelled
          ? `Subagent job ${jobId} cancelled.`
          : formatCancellationFailure(jobId, jobs),
        cancelled ? "info" : "warning",
      );
    },
  });

  const registerTool = async (): Promise<void> => {
    refreshSubagentTool = await registerSubagentTool(
      pi,
      limits,
      config.interactive,
      jobs,
      deliveries,
      pendingMessages,
      () => parentAgentActive,
      () => silentMode,
      () => shuttingDown,
      supportsAgentShellModelDiscovery(),
      supportsAgentShellInteractive(),
      config.roster.enabled,
    );
  };

  if (isAgentShellRuntimeInstalled()) {
    const modelDiscoverySupported = supportsAgentShellModelDiscovery();
    const interactiveSupported = supportsAgentShellInteractive();
    refreshSubagentTool = await registerSubagentTool(
      pi,
      limits,
      config.interactive,
      jobs,
      deliveries,
      pendingMessages,
      () => parentAgentActive,
      () => silentMode,
      () => shuttingDown,
      modelDiscoverySupported,
      interactiveSupported,
      config.roster.enabled,
    );

    if (
      !modelDiscoverySupported ||
      (config.interactive.enabled && !interactiveSupported)
    ) {
      pi.on("session_start", (_event, ctx) => {
        if (!ctx.hasUI) {
          return;
        }

        const notices = [];

        if (!modelDiscoverySupported) {
          notices.push([
            "Model discovery needs an updated AgentShell runtime.",
            `Run: ${setupCommand()}`,
            "Then restart Pi or run /reload.",
          ].join("\n"));
        }

        if (config.interactive.enabled && !interactiveSupported) {
          notices.push([
            "Interactive runs need an updated AgentShell runtime.",
            `Run: ${setupCommand()}`,
            "Then restart Pi or run /reload.",
          ].join("\n"));
        }

        ctx.ui.notify(
          notices.join("\n\n"),
          "warning",
        );
      });
    }

    return;
  }

  let setupAttempted = false;

  pi.on("session_start", async (_event, ctx) => {
    if (setupAttempted) {
      return;
    }

    setupAttempted = true;
    const command = setupCommand();

    if (!ctx.hasUI) {
      process.stderr.write(
        `AgentShell runtime is missing. Run: ${command}\n`,
      );
      return;
    }

    const uvCheck = await pi.exec(
      "uv",
      ["--version"],
      { timeout: 5_000 },
    );

    if (uvCheck.code !== 0) {
      ctx.ui.notify(
        [
          "`uv` is required to set up the AgentShell runtime.",
          "Install it from:",
          UV_INSTALL_URL,
          "Then restart Pi or run /reload.",
        ].join("\n"),
        "error",
      );
      return;
    }

    const confirmed = await ctx.ui.confirm(
      "Set up AgentShell runtime?",
      [
        "The subagent extension needs its locked Python dependency.",
        "",
        `Run: ${command}`,
        "",
        "This may download Python and the audited AgentShell dependency.",
      ].join("\n"),
    );

    if (!confirmed) {
      ctx.ui.notify(
        `AgentShell was not installed. Run manually:\n${command}`,
        "warning",
      );
      return;
    }

    ctx.ui.notify("Installing the AgentShell runtime...", "info");

    const setup = await pi.exec(
      "uv",
      [
        "sync",
        "--project",
        AGENT_SHELL_PROJECT_DIRECTORY,
        "--locked",
      ],
    );

    if (setup.code !== 0) {
      const diagnostic =
        setup.stderr.trim() ||
        setup.stdout.trim() ||
        `uv exited with code ${setup.code}`;

      ctx.ui.notify(
        `AgentShell setup failed:\n${diagnostic}`,
        "error",
      );
      return;
    }

    try {
      await registerTool();
      ctx.ui.notify("The subagent tool is ready.", "info");
    } catch (error) {
      const message =
        error instanceof Error ? error.message : String(error);

      ctx.ui.notify(
        `AgentShell was installed but could not start:\n${message}`,
        "error",
      );
    }
  });
}