/**
 * Mode Switcher — unified chat/plan/execute mode manager.
 *
 * Single shortcut (ctrl+alt+m) opens a picker to switch between:
 *   Chat     — read-only conversational mode
 *   Plan     — read-only explore-and-plan mode (with plan file management)
 *   Normal   — restore normal agent behavior
 *
 * Replaces the separate chat-mode and plan-mode extensions. Publishes both
 * __chatMode and __planMode globals so chat-input and footer segments continue
 * to work without changes.
 *
 * Backward-compatible commands: /chat, /plan, /mode
 * Backward-compatible flags:    --chat, --plan
 */

import type {
  ExtensionAPI,
  ExtensionContext,
  BeforeAgentStartEvent,
  BeforeAgentStartEventResult,
  ToolCallEvent,
  ToolCallEventResult,
  ToolResultEvent,
  AgentEndEvent,
  AgentBeforeSettleEvent,
} from "@earendil-works/pi-coding-agent";
import { DynamicBorder } from "@earendil-works/pi-coding-agent";
import type { Model } from "@earendil-works/pi-ai";
import { Type } from "typebox";
import { Container, Input, Text, Spacer, type KeyId, type Component, type SelectItem } from "@earendil-works/pi-tui";
import { join } from "node:path";
import { existsSync, writeFileSync, readFileSync, unlinkSync } from "node:fs";

import {
  PLAN_MODE_TOOLS,
  CHAT_MODE_TOOLS,
  PLAN_MODE_PROMPT,
  CHAT_MODE_PROMPT,
  PLAN_FILE_PREFIX,
  PLAN_DIR,
  buildExecutePrompt,
  buildRefinePrompt,
  USER_CONFIG,
  MODE_REGISTRY,
  getModeDefinition,
  executeHandoffAllowed,
  withPlanComplete,
  planCompleteAllowed,
  LAZY_TOOLS,
  LOADER_TOOL_NAME,
  applyLazyPolicy,
  lazyToolsToEnable,
} from "./config.js";
import {
  bashBlockReason,
  extractPlanText,
  isPlanLike,
  ensurePlanDir,
  titleFromFilename,
  listPlanFiles,
  sanitizePlanName,
  extractTextFromMessage,
  applyLabelColor,
  createLatestModelQueue,
} from "./utils.js";
import {
  getMode,
  getRefining,
  setRefining,
  getActivePlanFile,
  setActivePlanFile,
  getRestoringModel,
  setRestoringModel,
  transition,
  enterPlanWithFile,
  restore,
  resetState,
  getRefineCount,
  incrementRefineCount,
} from "./state.js";
import { showSelectMenu } from "./menus.js";

export default function modeSwitcher(pi: ExtensionAPI) {
  // ─── Lazy tool loader ──────────────────────────────────────────────────────
  // Registered only when lazy tools are configured. Tools in LAZY_TOOLS are
  // withheld from the active set at rest (see applyLazyPolicy) and enabled on
  // demand here, bounded by the current mode's tool policy.
  if (LAZY_TOOLS.size > 0) {
    pi.registerTool({
      name: LOADER_TOOL_NAME,
      label: "Load Tools",
      description:
        `Enable on-demand tools that are withheld from the prompt until needed. ` +
        `Available lazy tools: ${[...LAZY_TOOLS].join(", ")}. ` +
        `Call this before using one of them. Activation is bounded by the current mode's tool policy, ` +
        `so read-only modes cannot enable tools they do not permit.`,
      parameters: Type.Object({
        tools: Type.Optional(
          Type.Array(Type.String(), {
            description: "Tool names to enable. Omit to enable all lazy tools the current mode allows.",
          }),
        ),
      }),
      async execute(_toolCallId, params) {
        const mode = getMode();
        const modeDef = getModeDefinition(mode);
        const modeAllowed = modeDef?.tools ?? null; // null = mode inherits all tools
        const current = pi.getActiveTools();
        const enable = lazyToolsToEnable(params.tools as string[] | undefined, current, modeAllowed);
        if (enable.length > 0) pi.setActiveTools([...current, ...enable]);
        const text =
          enable.length > 0
            ? `Enabled: ${enable.join(", ")}`
            : "No tools enabled (already active, unknown, or not allowed in the current mode).";
        return { content: [{ type: "text", text }], details: { enabled: enable } };
      },
    });
  }

  // ─── Saved tool list for restoring ────────────────────────────────────────
  let savedToolNames: string[] | null = null;
  // ─── Saved model for restoring ─────────────────────────────────────────────
  let savedModel: Model<any> | null = null;
  const queueModelChange = createLatestModelQueue();
  // Consecutive restore failures for the same model, so an un-authable restore point is
  // retried once and then released instead of blocking future captures forever.
  let restoreAttempts = 0;

  function saveAndSetActiveTools(toolNames: string[], resting = false): void {
    if (resting) {
      // OFF is the resting baseline: drop any snapshot from the mode being exited, otherwise a
      // later tool-less mode would restore the pre-OFF set and widen past the pinned baseline.
      savedToolNames = null;
    } else if (savedToolNames === null) {
      // Snapshot what is actually active, not every registered tool: entering a mode must
      // never widen the set the user (or Pi settings) narrowed before it.
      savedToolNames = pi.getActiveTools();
    }
    pi.setActiveTools(applyLazyPolicy(toolNames));
  }

  /** Plan-mode tool list: honors a modes.plan.tools override, not just the shared default. */
  function planModeTools(): string[] {
    return getModeDefinition("plan")?.tools ?? PLAN_MODE_TOOLS;
  }

  /**
   * Pi's setActiveTools() silently drops names it does not know, so a typo in
   * modes.<name>.tools quietly shrinks the mode. Say so once per name per session.
   */
  const warnedTools = new Set<string>();
  function warnUnknownTools(ctx: ExtensionContext, mode: string, tools?: string[]): void {
    if (!ctx.hasUI || !tools) return;
    const known = new Set(pi.getAllTools().map((t) => t.name));
    for (const name of tools) {
      if (known.has(name) || warnedTools.has(name)) continue;
      warnedTools.add(name);
      ctx.ui.notify(`[mode-switcher] modes.${mode}.tools: unknown tool "${name}" is ignored by Pi`, "warning");
    }
  }

  /** Queue model changes so a superseded async switch cannot finish after its restore. */
  function setQueuedModel(ctx: ExtensionContext, model: Model<any>, label: string): Promise<boolean | undefined> {
    return queueModelChange(async () => {
      const success = await pi.setModel(model);
      if (!success && ctx.hasUI) ctx.ui.notify(`[mode-switcher] No API key for model ${label}`, "error");
      updateStatus(ctx);
      return success;
    });
  }

  async function switchToModel(ctx: ExtensionContext, modelRef: { provider: string; id: string }): Promise<void> {
    const resolved = ctx.modelRegistry.find(modelRef.provider, modelRef.id);
    if (!resolved) {
      if (ctx.hasUI) ctx.ui.notify(`[mode-switcher] Model not found: ${modelRef.provider}/${modelRef.id}`, "warning");
      return;
    }
    await setQueuedModel(ctx, resolved, `${modelRef.provider}/${modelRef.id}`);
  }

  /**
   * Switch to the mode's configured model (if any). Awaiting this guarantees the
   * model switch has actually landed before any dependent action (e.g. starting
   * execution). Callers that don't need that guarantee (plain mode entry) can
   * ignore the returned promise — the footer's model segment (which reads
   * ctx.model live on every render) still gets refreshed via updateStatus(ctx)
   * once the switch resolves.
   */
  async function applyModeModel(ctx: ExtensionContext, modeDef: { model?: { provider: string; id: string } } | undefined): Promise<void> {
    const modelRef = modeDef?.model;
    if (!modelRef) return;
    const current = ctx.model ?? null;
    const alreadyOnModeModel = current !== null && current.provider === modelRef.provider && current.id === modelRef.id;
    // Capture the restore point once per mode, and never when the session is already on
    // the mode's own model (after /reload or /resume inside a mode) or when a restore
    // point is already persisted — otherwise the real previous model is lost.
    if (savedModel === null && getRestoringModel() === null && !alreadyOnModeModel) {
      rememberRestoreModel(current);
    }
    await switchToModel(ctx, modelRef);
  }

  /** Record the model to return to, in memory and in the session blob. */
  function rememberRestoreModel(model: Model<any> | null): void {
    savedModel = model;
    if (model) restoreAttempts = 0;
    setRestoringModel(model ? { provider: model.provider, id: model.id } : null, pi);
  }

  /** Restore whatever model was active before applyModeModel last changed it. */
  function restoreModelIfSaved(ctx: ExtensionContext): void {
    const ref = savedModel ?? getRestoringModel();
    const toRestore = ref ? (savedModel ?? ctx.modelRegistry.find(ref.provider, ref.id)) : null;
    rememberRestoreModel(null);
    // The model left the registry (provider removed): drop the point so the next entry captures a fresh one.
    if (!toRestore) return;
    void setQueuedModel(ctx, toRestore, `${toRestore.provider}/${toRestore.id}`).then((outcome) => {
      // false = the model exists but has no auth yet (fixed by /login): keep it for one retry.
      // undefined = superseded by a newer switch, which now owns the restore.
      if (outcome === false && ref && restoreAttempts < 1) {
        restoreAttempts++;
        setRestoringModel({ provider: ref.provider, id: ref.id }, pi);
      } else {
        restoreAttempts = 0;
      }
    });
  }

  /**
   * OFF mode's model handling. If OFF has its own configured model, apply it
   * deterministically — OFF is the resting baseline, so drop any pending restore
   * so future mode switches capture OFF's model as their restore point. If OFF
   * has no model configured, fall back to restoring whatever model was active
   * before the last mode switch (original behavior).
   */
  function applyOffModel(ctx: ExtensionContext): void {
    const offModeDef = getModeDefinition("off");
    if (offModeDef?.model) {
      rememberRestoreModel(null);
      void switchToModel(ctx, offModeDef.model);
    } else {
      restoreModelIfSaved(ctx);
    }
  }

  /** Filter out plan_complete — only available in execute mode. */
  function toolsWithoutPlanComplete(names: string[]): string[] {
    return names.filter((n) => n !== "plan_complete");
  }

  function restoreAllTools(): void {
    const names = savedToolNames ?? pi.getActiveTools();
    pi.setActiveTools(applyLazyPolicy(toolsWithoutPlanComplete(names)));
    savedToolNames = null;
  }

  // ─── Plan file helpers ─────────────────────────────────────────────────────

  function getPlanFilePath(): string | null {
    const file = getActivePlanFile();
    if (!file) return null;
    return join(process.cwd(), PLAN_DIR, file);
  }

  function getPlanDisplayTitle(): string | null {
    const file = getActivePlanFile();
    if (!file) return null;
    const filePath = getPlanFilePath();
    if (filePath && existsSync(filePath)) {
      const heading = readFileSync(filePath, "utf-8").match(/^# Plan:\s*(.+)$/m);
      if (heading) return heading[1].trim();
    }
    return titleFromFilename(file);
  }

  // ─── UI helpers ────────────────────────────────────────────────────────────

  function updateStatus(ctx: ExtensionContext): void {
    const mode = getMode();
    const modeDef = getModeDefinition(mode);
    warnUnknownTools(ctx, mode, modeDef?.tools);

    if (!ctx.hasUI) return;

    // Chat widget
    if (mode === "chat") {
      const widgetColor = modeDef?.labels?.widgetColor ?? USER_CONFIG.labels.chat.widgetColor;
      const widgetText = modeDef?.labels?.widget ?? USER_CONFIG.labels.chat.widget;
      ctx.ui.setWidget(
        "chat-mode",
        USER_CONFIG.ui.hideWidget ? undefined : [applyLabelColor(ctx.ui.theme, widgetColor, widgetText)],
      );
    } else {
      ctx.ui.setWidget("chat-mode", undefined);
    }

    // Plan/execute widget
    if (mode === "plan") {
      const title = getPlanDisplayTitle();
      const widgetWithTitle = modeDef?.labels?.widgetWithTitle ?? USER_CONFIG.labels.plan.widgetWithTitle;
      const widgetPlain = modeDef?.labels?.widget ?? USER_CONFIG.labels.plan.widget;
      const widgetText = title ? widgetWithTitle.replace("{title}", title) : widgetPlain;
      const widgetColor = modeDef?.labels?.widgetColor ?? USER_CONFIG.labels.plan.widgetColor;
      ctx.ui.setWidget(
        "plan-mode",
        USER_CONFIG.ui.hideWidget ? undefined : [applyLabelColor(ctx.ui.theme, widgetColor, widgetText)],
      );
    } else if (mode === "execute") {
      const title = getPlanDisplayTitle();
      const widgetWithTitle = modeDef?.labels?.widgetWithTitle ?? USER_CONFIG.labels.execute.widgetWithTitle;
      const widgetPlain = modeDef?.labels?.widget ?? USER_CONFIG.labels.execute.widget;
      const widgetText = title ? widgetWithTitle.replace("{title}", title) : widgetPlain;
      const widgetColor = modeDef?.labels?.widgetColor ?? USER_CONFIG.labels.execute.widgetColor;
      ctx.ui.setWidget(
        "plan-mode",
        USER_CONFIG.ui.hideWidget ? undefined : [applyLabelColor(ctx.ui.theme, widgetColor, widgetText)],
      );
    } else {
      ctx.ui.setWidget("plan-mode", undefined);
    }

    // Custom mode widget (any registered mode other than off/chat/plan/execute)
    if (mode !== "off" && mode !== "chat" && mode !== "plan" && mode !== "execute") {
      const widgetText = modeDef?.labels?.widget || `✓ ${mode} mode`;
      const widgetColor = modeDef?.labels?.widgetColor || "accent";
      ctx.ui.setWidget(
        "custom-mode",
        USER_CONFIG.ui.hideWidget ? undefined : [applyLabelColor(ctx.ui.theme, widgetColor, widgetText)],
      );
    } else {
      ctx.ui.setWidget("custom-mode", undefined);
    }

    // Emit legacy events for any other extensions listening
    pi.events.emit("chat-mode:state", { mode: mode === "chat" ? "chat" : "off" });
    pi.events.emit("plan-mode:state", { mode: mode === "chat" ? "off" : mode });
  }

  // ─── State transitions ─────────────────────────────────────────────────────

  function enterChatMode(ctx: ExtensionContext): void {
    transition("chat", pi);
    const modeDef = getModeDefinition("chat");
    saveAndSetActiveTools(modeDef?.tools ?? CHAT_MODE_TOOLS);
    applyModeModel(ctx, modeDef);
    updateStatus(ctx);
    if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
      const notify = modeDef?.labels?.notify ?? USER_CONFIG.labels.chat.notify;
      const notifyType = (modeDef?.labels?.notifyType ?? USER_CONFIG.labels.chat.notifyType) as "info" | "warning" | "error";
      ctx.ui.notify(notify, notifyType);
    }
  }

  function enterPlanMode(ctx: ExtensionContext): void {
    transition("plan", pi);
    const modeDef = getModeDefinition("plan");
    saveAndSetActiveTools(planModeTools());
    applyModeModel(ctx, modeDef);
    updateStatus(ctx);
    if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
      const notify = modeDef?.labels?.notify ?? USER_CONFIG.labels.plan.notify;
      const notifyType = (modeDef?.labels?.notifyType ?? USER_CONFIG.labels.plan.notifyType) as "info" | "warning" | "error";
      ctx.ui.notify(notify, notifyType);
    }
  }

  async function enterExecuteMode(ctx: ExtensionContext): Promise<void> {
    transition("execute", pi);
    const baseNames = savedToolNames ?? pi.getActiveTools();
    pi.setActiveTools(applyLazyPolicy([...toolsWithoutPlanComplete(baseNames), "plan_complete"]));
    savedToolNames = null;
    const modeDef = getModeDefinition("execute");
    await applyModeModel(ctx, modeDef);
    updateStatus(ctx);
    const title = getPlanDisplayTitle();
    const notifyWithTitle = modeDef?.labels?.notifyWithTitle ?? USER_CONFIG.labels.execute.notifyWithTitle;
    const notifyPlain = modeDef?.labels?.notify ?? USER_CONFIG.labels.execute.notify;
    const notifyText = title ? notifyWithTitle.replace("{title}", title) : notifyPlain;
    const notifyType = (modeDef?.labels?.notifyType ?? USER_CONFIG.labels.execute.notifyType) as "info" | "warning" | "error";
    if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
      ctx.ui.notify(notifyText, notifyType);
    }
  }

  /**
   * Exit whatever mode is currently active.
   * Pass silent=true when transitioning between modes (not going all the way to off)
   * to suppress the "Mode OFF" notification.
   */
  function enterOffMode(ctx: ExtensionContext, message?: string, silent = false): void {
    transition("off", pi);
    const offModeDef = getModeDefinition("off");
    // OFF gates tools only when the user declares modes.off.tools; otherwise it restores
    // whatever was active before the mode.
    if (offModeDef?.tools) {
      saveAndSetActiveTools(withPlanComplete("off", offModeDef.tools), true);
    } else {
      restoreAllTools();
    }
    applyOffModel(ctx);
    updateStatus(ctx);
    if (!silent && ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
      const modeDef = getModeDefinition("off");
      ctx.ui.notify(
        message ?? modeDef?.labels?.notify ?? USER_CONFIG.labels.off.notify,
        message ? "info" : ((modeDef?.labels?.notifyType ?? USER_CONFIG.labels.off.notifyType) as "info" | "warning" | "error"),
      );
    }
  }

  /** Generic entry for user-defined (non-built-in) modes registered via config. */
  function enterCustomMode(ctx: ExtensionContext, name: string): void {
    const modeDef = getModeDefinition(name);
    transition(name, pi);
    if (modeDef?.tools) {
      // Honor allowPlanComplete for custom modes: append plan_complete when set.
      saveAndSetActiveTools(withPlanComplete(name, modeDef.tools));
    } else {
      restoreAllTools();
      // Inherited (unlisted) tool set: allowPlanComplete still needs the tool active,
      // otherwise the mode can never finish through the completion path.
      if (modeDef?.allowPlanComplete) pi.setActiveTools(withPlanComplete(name, pi.getActiveTools()));
    }
    applyModeModel(ctx, modeDef);
    updateStatus(ctx);
    if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
      const notifyText = modeDef?.labels?.notify ?? `✓ ${name} mode ON`;
      const notifyType = (modeDef?.labels?.notifyType as "info" | "warning" | "error") ?? "info";
      ctx.ui.notify(notifyText, notifyType);
    }
  }

  /** Cleanup plan file on plan_complete, if configured. */
  function cleanupPlanFile(): void {
    if (!USER_CONFIG.cleanup.cleanupOnComplete) return;
    const filePath = getPlanFilePath();
    if (filePath && existsSync(filePath)) unlinkSync(filePath);
    setActivePlanFile(null, pi);
  }

  // ─── Mode picker ───────────────────────────────────────────────────────────

  /** Open the mode picker. Handles all mode transitions including cross-mode switches. */
  async function openModePicker(ctx: ExtensionContext): Promise<void> {
    const current = getMode();
    const planFiles = listPlanFiles();

    // Build mode list dynamically from registry (skip invisible/disabled modes)
    const registryModes = Array.from(MODE_REGISTRY.entries())
      .filter(([_, def]) => def.visible !== false && def.enabled !== false)
      .map(([name, def]) => {
        const isActive = current === name;
        const label = def.labels?.widget || name;
        return { value: name, label: isActive ? `${label.replace(/ active$/, '')}   (active)` : label };
      });
    
    // Always include execute options for plan files — unless the active mode
    // disables the execute handoff via modes.<name>.allowExecute: false.
    const executeItems = executeHandoffAllowed(current)
      ? planFiles.map(f => ({ value: `execute:${f.name}`, label: `Execute: ${f.title}` }))
      : [];
    
    const items: SelectItem[] = [...registryModes, ...executeItems];

    const choice = await showSelectMenu(ctx, "Select mode", items);
    if (!choice) return; // cancelled

    if (choice === "off") {
      if (current === "off") return;
      enterOffMode(ctx);
    } else if (choice === "chat") {
      if (current === "chat") return;
      if (current !== "off") enterOffMode(ctx, undefined, true);
      enterChatMode(ctx);
    } else if (choice === "plan") {
      if (current === "plan" || current === "execute") return;
      if (current !== "off") enterOffMode(ctx, undefined, true);
      await promptNameAndEnterPlanMode(ctx);
    } else if (choice.startsWith("execute:")) {
      if (!executeHandoffAllowed(current)) {
        if (ctx.hasUI) ctx.ui.notify(`Plan execution is not available from ${current} mode.`, "warning");
        return;
      }
      const filename = choice.slice("execute:".length);
      if (current !== "off") enterOffMode(ctx, undefined, true);
      enterPlanWithFile(filename, pi);
      saveAndSetActiveTools(planModeTools());
      await enterExecuteMode(ctx);
      pi.sendUserMessage("Execute the plan steps now.", { deliverAs: "followUp" });
    } else if (getModeDefinition(choice)) {
      // Custom mode registered via config
      if (current === choice) return;
      if (current !== "off") enterOffMode(ctx, undefined, true);
      enterCustomMode(ctx, choice);
    }
  }

  // ─── Re-derive state from entries on the current branch ──────────────────────

  async function syncStateFromBranch(ctx: ExtensionContext, autoResume = false): Promise<void> {
    const branch = ctx.sessionManager.getBranch();
    const restored = restore(branch);

    if (!restored) {
      resetState();
      restoreAllTools();
      applyOffModel(ctx);
      updateStatus(ctx);
      return;
    }

    // Check if active plan file was deleted externally
    const filePath = getPlanFilePath();
    if (getActivePlanFile() && filePath && !existsSync(filePath)) {
      if (ctx.hasUI) ctx.ui.notify(`Plan file "${getActivePlanFile()}" not found — disabling plan mode.`, "warning");
      enterOffMode(ctx);
      return;
    }

    const mode = getMode();
    const modeDef = getModeDefinition(mode);

    if (modeDef?.tools) {
      // Honor allowPlanComplete here too, otherwise a resume drops the tool a mode was given.
      // OFF's declared list is the resting baseline, so it must not keep a stale snapshot.
      saveAndSetActiveTools(withPlanComplete(mode, modeDef.tools), mode === "off");
    } else if (mode === "execute") {
      const allNames = pi.getActiveTools();
      pi.setActiveTools(applyLazyPolicy([...toolsWithoutPlanComplete(allNames), "plan_complete"]));
    } else {
      // Default: the active set minus plan_complete, re-added when the mode allows it.
      const allNames = pi.getActiveTools();
      pi.setActiveTools(applyLazyPolicy(withPlanComplete(mode, toolsWithoutPlanComplete(allNames))));
    }

    // Apply the mode's model BEFORE sending any resume follow-up message, so
    // execution resumes on the correct model rather than racing the switch.
    await applyModeModel(ctx, modeDef);
    updateStatus(ctx);

    if (mode === "execute" && autoResume) {
      const title = getPlanDisplayTitle();
      const msg = title
        ? `Resuming execution of plan: "${title}". Review the plan in context and continue from where execution left off.`
        : "Resuming plan execution. Review the plan in context and continue from where execution left off.";
      pi.sendUserMessage(msg, { deliverAs: "followUp" });
    }
  }

  // ─── Event: session_start ──────────────────────────────────────────────────────

  pi.on("session_start", async (event, ctx) => {
    if (event.reason === "startup") {
      if (pi.getFlag("chat") === true && getMode() === "off") {
        enterChatMode(ctx);
        return;
      }
      if (pi.getFlag("plan") === true && getMode() === "off") {
        enterPlanMode(ctx);
        return;
      }
    }

    await syncStateFromBranch(ctx, event.reason === "resume" || event.reason === "fork");
  });

  // ─── Event: before_agent_start ──────────────────────────────────────────

  pi.on("before_agent_start", (event: BeforeAgentStartEvent): BeforeAgentStartEventResult => {
    const mode = getMode();
    const modeDef = getModeDefinition(mode);
    
    if (!modeDef) return {};
    
    // Special handling for plan mode with refining
    if (mode === "plan" && getRefining()) {
      const filePath = getPlanFilePath();
      if (filePath && existsSync(filePath)) {
        const planContent = readFileSync(filePath, "utf-8");
        return { systemPrompt: event.systemPrompt + "\n\n" + buildRefinePrompt(planContent) };
      }
    }
    
    // Special handling for execute mode with plan content
    if (mode === "execute") {
      const filePath = getPlanFilePath();
      if (!filePath || !existsSync(filePath)) return {};
      const planContent = readFileSync(filePath, "utf-8");
      const customTemplate = modeDef.prompt;
      if (typeof customTemplate === "string" && customTemplate.trim().length > 0) {
        // User-configured template via modes.execute.prompt — {plan} is replaced with the plan file content.
        const filled = customTemplate.includes("{plan}")
          ? customTemplate.replace("{plan}", planContent)
          : `${customTemplate}\n\nPlan:\n${planContent}`;
        return { systemPrompt: event.systemPrompt + "\n\n" + filled };
      }
      return { systemPrompt: event.systemPrompt + "\n\n" + buildExecutePrompt(planContent) };
    }
    
    // Generic mode prompt from registry
    if (modeDef.prompt) {
      const promptText = typeof modeDef.prompt === "string" 
        ? modeDef.prompt 
        : modeDef.prompt();
      if (promptText) {
        return { systemPrompt: event.systemPrompt + "\n\n" + promptText };
      }
    }
    
    return {};
  });

  // ─── Event: tool_call ─────────────────────────────────────────────────────

  pi.on("tool_call", (event: ToolCallEvent): ToolCallEventResult => {
    // plan_complete is only callable in execute mode
    const mode = getMode();
    const modeDef = getModeDefinition(mode);
    if (event.toolName === "plan_complete" && !planCompleteAllowed(mode)) {
      return {
        block: true,
        reason: `[mode-switcher] plan_complete is only available in execute mode or a mode with allowPlanComplete. Current mode: ${mode}`,
      };
    }

    // Bash safety check using registry patterns
    if (event.toolName !== "bash") return {};
    if (!modeDef?.safePatterns && !modeDef?.destructivePatterns) return {};

    const command = event.input.command as string;
    const reason = bashBlockReason(command, modeDef.safePatterns, modeDef.destructivePatterns);
    if (reason) {
      return { block: true, reason: `[mode-switcher] Command blocked — ${reason} [${mode} mode]` };
    }

    return {};
  });

  // ─── Tool: plan_complete ────────────────────────────────────────────────────

  pi.registerTool({
    name: "plan_complete",
    label: "Plan Complete",
    description:
      "Signal that all plan steps have been executed. ONLY callable in EXECUTE mode or in a mode configured with allowPlanComplete. Call this once after finishing the final step. This exits the mode. Do NOT call this outside those modes.",
    parameters: Type.Object({}),
    async execute(_toolCallId, _params, _signal, _onUpdate, ctx) {
      if (!planCompleteAllowed(getMode())) {
        return {
          content: [{ type: "text", text: "plan_complete is only available in execute mode or a mode with allowPlanComplete." }],
          details: undefined,
        };
      }
      cleanupPlanFile();
      enterOffMode(ctx, "Plan implemented. Plan mode OFF.");
      return { content: [{ type: "text", text: "Plan mode exited." }], details: undefined };
    },
  });

  // ─── Event: tool_result (process plan_complete results) ───────────────────

  pi.on("tool_result", (event: ToolResultEvent, ctx: ExtensionContext) => {
    if (event.toolName !== "plan_complete") return;
    // The tool body has usually already returned to OFF, which makes this a no-op; the check
    // stays mode-based (not execute-only) so a custom allowPlanComplete mode is handled alike.
    if (getMode() !== "execute" && !planCompleteAllowed(getMode())) return;
    cleanupPlanFile();
    enterOffMode(ctx, "Plan implemented. Plan mode OFF.");
  });

  // ─── Event: agent_before_settle (auto-exit execute when the run is truly done) ───

  // Pi 0.87 boundary: fires after any automatic retry, compaction-and-retry, or queued
  // follow-up work, so execute mode is not dropped mid-plan at an intermediate agent_end.
  pi.on("agent_before_settle", async (event: AgentBeforeSettleEvent, ctx: ExtensionContext) => {
    // If plan_complete was never called, exit execute mode automatically once Pi has fully
    // settled. (If it was called, mode is already "off" here — this check is a no-op.)
    // "aborted" (ESC) and "error" (provider failure) are pauses, not finished executions;
    // only a completed run exits execute mode.
    if (getMode() === "execute" && event.outcome === "completed") {
      enterOffMode(ctx, "Execution complete. Mode OFF.");
    }
  });

  // ─── Event: agent_end (extract plan text in PLAN mode) ───

  pi.on("agent_end", async (event: AgentEndEvent, ctx: ExtensionContext) => {
    if (getMode() !== "plan") return;
    if (!ctx.hasUI) return;

    const lastAssistant = [...event.messages].reverse().find((m) => m.role === "assistant");
    if (!lastAssistant) return;

    const text = extractTextFromMessage(lastAssistant as unknown as Record<string, unknown>);
    if (!text) return;

    const planText = extractPlanText(text);
    setRefining(false);

    if (planText) {
      const planDir = ensurePlanDir();
      const activeFile = getActivePlanFile();
      let filename = activeFile;
      if (!filename) {
        const ts = new Date().toISOString().replace(/[T:]/g, "-").slice(0, 16);
        filename = `${PLAN_FILE_PREFIX}${ts}.md`;
        setActivePlanFile(filename, pi);
      }
      const filePath = join(planDir, filename);
      const title = titleFromFilename(filename);
      const content = `# Plan: ${title}\n\n${planText.replace(/^\s*#{1,6}\s*Plan:[^\n]*\n?/, "").trimStart()}\n`;
      writeFileSync(filePath, content, "utf-8");
      updateStatus(ctx);

      await showPlanMenu(ctx);
    } else if (isPlanLike(text)) {
      updateStatus(ctx);
      await showPlanMenu(ctx);
    }
  });

  // ─── Event: session_tree ──────────────────────────────────────────────────

  pi.on("session_tree", async (_event, ctx) => {
    await syncStateFromBranch(ctx);
  });

  // ─── Prompt user for optional plan name ────────────────────────────────────

  async function promptNameAndEnterPlanMode(ctx: ExtensionContext): Promise<boolean> {
    const nameInput = await ctx.ui.custom<string | undefined>((tui, theme, kb, done) => {
      const input = new Input();
      input.onSubmit = (value) => done(value);

      const border = new DynamicBorder((s: string) => theme.fg("border", s));
      const label = new Text(theme.fg("text", "Plan name ") + theme.fg("dim", "(optional, leave empty for timestamp)"), 1, 0);
      const indent = 1;

      const container = new Container();
      container.addChild(border);
      container.addChild(new Spacer());
      container.addChild(label);
      container.addChild(new Spacer());
      const indentedInput: Component = {
        render: (w: number) => input.render(w - indent).map((line) => " ".repeat(indent) + line),
        invalidate: () => input.invalidate(),
      };
      container.addChild(indentedInput);
      container.addChild(new Spacer());
      container.addChild(new Text(theme.fg("dim", "enter") + theme.fg("muted", " to submit") + theme.fg("dim", " • ") + theme.fg("dim", "esc") + theme.fg("muted", " to cancel"), 1, 0));
      container.addChild(new Spacer());
      container.addChild(border);

      input.focused = true;

      return {
        render(width: number) { return container.render(width); },
        invalidate() { container.invalidate(); },
        handleInput(data: string) {
          if (data === "\x1b") { done(undefined); return; } // esc
          input.handleInput(data);
          tui.requestRender();
        },
      };
    });

    if (nameInput === undefined) return false; // cancelled

    if (nameInput.trim()) {
      const sanitized = sanitizePlanName(nameInput.trim());
      if (!sanitized) {
        if (ctx.hasUI) ctx.ui.notify("Invalid plan name. Use letters, numbers, hyphens, underscores, spaces, and dots only.", "warning");
        return false;
      }
      const filename = `${PLAN_FILE_PREFIX}${sanitized}.md`;
      enterPlanWithFile(filename, pi);
      saveAndSetActiveTools(planModeTools());
      updateStatus(ctx);
      const createTitle = titleFromFilename(filename);
      if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
        ctx.ui.notify(USER_CONFIG.labels.plan.notifyWithTitle.replace("{title}", createTitle), USER_CONFIG.labels.plan.notifyType as "info" | "warning" | "error");
      }
    } else {
      enterPlanMode(ctx);
    }
    return true;
  }

  /** Show the plan action menu: Execute, Refine, Save & Exit, Discard & Exit. */
  async function showPlanMenu(ctx: ExtensionContext): Promise<void> {
    while (true) {
      if (getMode() !== "plan") return;

      const refineCount = getRefineCount();
      const refineLabel = refineCount > 0
        ? `Refine  (${refineCount} ${refineCount === 1 ? "cycle" : "cycles"} so far)`
        : "Refine";
      const options: SelectItem[] = [
        { value: "execute", label: "Execute" },
        { value: "refine", label: refineLabel },
        { value: "save", label: "Save & Exit" },
        { value: "discard", label: "Discard & Exit" },
      ];

      const choice = await showSelectMenu(ctx, "How'd you like to proceed?", options);

      if (choice === "execute") {
        await enterExecuteMode(ctx);
        pi.sendUserMessage("Execute the plan steps now.", { deliverAs: "followUp" });
        return;
      } else if (choice === "refine") {
        incrementRefineCount();
        setRefining(true);
        updateStatus(ctx);
        return;
      } else if (choice === "save") {
        enterOffMode(ctx, "Plan saved. Plan mode OFF.");
        return;
      } else if (choice === "discard") {
        const confirmOptions: SelectItem[] = [
          { value: "yes", label: "Yes, I am sure" },
          { value: "cancel", label: "Cancel" },
        ];
        const confirmed = await showSelectMenu(ctx, "Are you sure?", confirmOptions, {
          bold: false,
          dimSuffix: " (this will delete the plan file)",
        });

        if (confirmed === "yes") {
          const filePath = getPlanFilePath();
          if (filePath && existsSync(filePath)) unlinkSync(filePath);
          setActivePlanFile(null, pi);
          enterOffMode(ctx, "Plan discarded. Plan mode OFF.");
          return;
        }
        // Cancel — re-show plan menu
      } else {
        // Escaped/cancelled — re-show plan menu (noop, loop continues)
      }
    }
  }

  /** Load an existing plan file and show the action menu. */
  function loadPlanAndShowMenu(ctx: ExtensionContext, filename: string, displayName: string): void {
    enterPlanWithFile(filename, pi);
    saveAndSetActiveTools(planModeTools());
    updateStatus(ctx);
    if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
      ctx.ui.notify(USER_CONFIG.labels.plan.notifyLoaded.replace("{title}", displayName), "info");
    }

    const filePath = join(process.cwd(), PLAN_DIR, filename);
    if (existsSync(filePath)) {
      const planContent = readFileSync(filePath, "utf-8");
      pi.sendMessage({
        customType: "plan-mode",
        content: planContent,
        display: true,
        details: { title: `Active plan: ${displayName}` },
      });
    }
  }

  // ─── Message renderer for plan messages ────────────────────────────────────

  pi.registerMessageRenderer("plan-mode", (message, _options, theme) => {
    const border = new DynamicBorder((s: string) => theme.fg("border", s));
    const container = new Container();
    container.addChild(border);
    container.addChild(new Text(theme.fg("text", "📄 " + (message.content as string).split("\n")[0])));
    const body = (message.content as string).split("\n").slice(1).join("\n");
    if (body) {
      container.addChild(new Spacer());
      container.addChild(new Text(body));
    }
    container.addChild(border);
    return container;
  });

  // ─── Command: /mode [chat|plan|off|normal|<custom>] ──────────────────────

  /** Names /mode accepts as an argument: the built-ins plus every registered custom mode. */
  function modeArgumentNames(): string[] {
    return ["off", "normal", "chat", "plan", ...Array.from(MODE_REGISTRY.keys()).filter((n) => n !== "off" && n !== "chat" && n !== "plan" && n !== "execute")];
  }

  pi.registerCommand("mode", {
    description: "Mode switcher: /mode (open picker) · /mode chat · /mode plan · /mode normal · /mode <custom>",
    handler: async (args: string, ctx) => {
      const input = args.trim().toLowerCase();

      if (!input) {
        await openModePicker(ctx);
        return;
      }

      if (input === "off" || input === "normal") {
        if (getMode() === "off") {
          if (ctx.hasUI) ctx.ui.notify("Already in normal mode", "info");
          return;
        }
        enterOffMode(ctx);
        return;
      }

      if (input === "chat") {
        if (getMode() === "chat") {
          if (ctx.hasUI) ctx.ui.notify("Already in chat mode", "info");
          return;
        }
        if (getMode() !== "off") enterOffMode(ctx, undefined, true);
        enterChatMode(ctx);
        return;
      }

      if (input === "plan") {
        if (getMode() === "plan" || getMode() === "execute") {
          if (ctx.hasUI) ctx.ui.notify("Already in plan mode", "info");
          return;
        }
        if (getMode() !== "off") enterOffMode(ctx, undefined, true);
        await promptNameAndEnterPlanMode(ctx);
        return;
      }

      // Any other registered mode: config-declared custom modes (modes.<name>) and built-in
      // overrides, so /mode review works the same as picking it from the picker.
      if (getModeDefinition(input)) {
        if (input === "execute") {
          if (ctx.hasUI) ctx.ui.notify("Execute mode needs an active plan — use /execute", "info");
          return;
        }
        if (getModeDefinition(input)?.enabled === false) {
          if (ctx.hasUI) ctx.ui.notify(`Mode "${input}" is disabled (modes.${input}.enabled: false)`, "warning");
          return;
        }
        if (getMode() === input) {
          if (ctx.hasUI) ctx.ui.notify(`Already in ${input} mode`, "info");
          return;
        }
        if (getMode() !== "off") enterOffMode(ctx, undefined, true);
        enterCustomMode(ctx, input);
        return;
      }

      if (ctx.hasUI) ctx.ui.notify(`Unknown mode: "${input}". Use: ${modeArgumentNames().join(", ")}`, "warning");
    },
  });

  // ─── Command: /chat [off] (alias) ──────────────────────────────────────────

  pi.registerCommand("chat", {
    description: "Chat mode: /chat (toggle) · /chat off",
    handler: async (args: string, ctx) => {
      const input = args.trim().toLowerCase();

      if (input === "off") {
        if (getMode() === "off") {
          if (ctx.hasUI) ctx.ui.notify("Chat mode is already off", "info");
          return;
        }
        enterOffMode(ctx);
        return;
      }

      if (getMode() === "off") {
        enterChatMode(ctx);
      } else if (getMode() === "chat") {
        enterOffMode(ctx);
      } else {
        // in plan/execute — switch to chat
        enterOffMode(ctx, undefined, true);
        enterChatMode(ctx);
      }
    },
  });

  // ─── Command: /plan [off|<name>] (alias) ───────────────────────────────────

  pi.registerCommand("plan", {
    description: "Plan mode: /plan (toggle) · /plan <name> (load existing or create new) · /plan off",
    handler: async (args: string, ctx) => {
      const input = args.trim();

      if (!input) {
        const current = getMode();
        if (current === "off") {
          const files = listPlanFiles();
          if (files.length === 0) {
            await promptNameAndEnterPlanMode(ctx);
          } else {
            const selectItems: SelectItem[] = [
              { value: "__new__", label: "New plan…" },
              ...files.map((f) => ({ value: f.name, label: f.title })),
            ];
            const choice = await showSelectMenu(ctx, "Select plan", selectItems);
            if (!choice) return;
            if (choice === "__new__") {
              await promptNameAndEnterPlanMode(ctx);
            } else {
              loadPlanAndShowMenu(ctx, choice, titleFromFilename(choice));
              await showPlanMenu(ctx);
            }
          }
        } else if (current === "plan" || current === "execute") {
          enterOffMode(ctx);
        } else {
          // in chat — switch to plan
          enterOffMode(ctx, undefined, true);
          await promptNameAndEnterPlanMode(ctx);
        }
        return;
      }

      if (input.toLowerCase() === "off") {
        if (getMode() === "off") {
          if (ctx.hasUI) ctx.ui.notify("Plan mode is already off", "info");
          return;
        }
        enterOffMode(ctx);
        return;
      }

      // Treat input as plan name: load existing → show menu, or create new → plan mode
      const sanitized = sanitizePlanName(input);
      if (!sanitized) {
        if (ctx.hasUI) ctx.ui.notify("Invalid plan name. Use letters, numbers, hyphens, underscores, spaces, and dots only.", "warning");
        return;
      }
      const filename = `${PLAN_FILE_PREFIX}${sanitized}.md`;
      const filePath = join(process.cwd(), PLAN_DIR, filename);

      if (existsSync(filePath)) {
        if (getMode() !== "off") enterOffMode(ctx, undefined, true);
        loadPlanAndShowMenu(ctx, filename, titleFromFilename(filename));
        await showPlanMenu(ctx);
      } else {
        if (getMode() === "plan" || getMode() === "execute") {
          if (ctx.hasUI) ctx.ui.notify("Already in plan mode", "info");
          return;
        }
        if (getMode() !== "off") enterOffMode(ctx, undefined, true);
        enterPlanWithFile(filename, pi);
        saveAndSetActiveTools(planModeTools());
        updateStatus(ctx);
        if (ctx.hasUI && !USER_CONFIG.ui.hideNotify) {
          ctx.ui.notify(USER_CONFIG.labels.plan.notifyWithTitle.replace("{title}", titleFromFilename(filename)), USER_CONFIG.labels.plan.notifyType as "info" | "warning" | "error");
        }
      }
    },
  });

  // ─── Command: /execute [<name>] ────────────────────────────────────────

  pi.registerCommand("execute", {
    description: "Execute an existing plan: /execute (pick from list) · /execute <name>",
    handler: async (args: string, ctx) => {
      const input = args.trim();
      const current = getMode();

      if (current !== "off" && !executeHandoffAllowed(current)) {
        if (ctx.hasUI) ctx.ui.notify(`Plan execution is not available from ${current} mode.`, "warning");
        return;
      }

      async function startExecute(filename: string): Promise<void> {
        if (current !== "off") enterOffMode(ctx, undefined, true);
        enterPlanWithFile(filename, pi);
        saveAndSetActiveTools(planModeTools());
        await enterExecuteMode(ctx);
        pi.sendUserMessage("Execute the plan steps now.", { deliverAs: "followUp" });
      }

      if (!input) {
        if (current === "execute") {
          if (ctx.hasUI) ctx.ui.notify("Already executing a plan", "info");
          return;
        }
        if (current === "plan" && getActivePlanFile()) {
          await enterExecuteMode(ctx);
          pi.sendUserMessage("Execute the plan steps now.", { deliverAs: "followUp" });
          return;
        }
        const files = listPlanFiles();
        if (files.length === 0) {
          if (ctx.hasUI) ctx.ui.notify("No plan files found. Run /plan first.", "warning");
          return;
        }
        if (files.length === 1) {
          await startExecute(files[0].name);
          return;
        }
        const items: SelectItem[] = files.map(f => ({ value: f.name, label: f.title }));
        const choice = await showSelectMenu(ctx, "Select plan to execute", items);
        if (!choice) return;
        await startExecute(choice);
        return;
      }

      const sanitized = sanitizePlanName(input);
      if (!sanitized) {
        if (ctx.hasUI) ctx.ui.notify("Invalid plan name.", "warning");
        return;
      }
      const filename = `${PLAN_FILE_PREFIX}${sanitized}.md`;
      const filePath = join(process.cwd(), PLAN_DIR, filename);
      if (!existsSync(filePath)) {
        if (ctx.hasUI) ctx.ui.notify(`Plan "${input}" not found. Check /plan for available plans.`, "warning");
        return;
      }
      await startExecute(filename);
    },
  });

  // ─── Shortcut: cycle mode ──────────────────────────────────────────────────

  /**
   * Cycle through all enabled, visible modes in registry order (excluding "execute",
   * which requires an active plan and is only entered via the picker).
   */
  function cycleMode(ctx: ExtensionContext): void {
    const current = getMode();
    const order = Array.from(MODE_REGISTRY.entries())
      .filter(([name, def]) => name !== "execute" && def.visible !== false && def.enabled !== false)
      .map(([name]) => name);
    if (order.length === 0) return;

    const idx = order.indexOf(current);
    const next = idx === -1 ? order[0] : order[(idx + 1) % order.length];
    if (next === current) return;

    if (next === "off") {
      enterOffMode(ctx);
    } else if (next === "chat") {
      if (current !== "off") enterOffMode(ctx, undefined, true);
      enterChatMode(ctx);
    } else if (next === "plan") {
      if (current === "execute") return; // don't leave an active execution via cycling
      if (current !== "off") enterOffMode(ctx, undefined, true);
      enterPlanMode(ctx);
    } else {
      if (current !== "off") enterOffMode(ctx, undefined, true);
      enterCustomMode(ctx, next);
    }
  }

  pi.registerShortcut(USER_CONFIG.shortcuts.cycleMode as KeyId, {
    description: "Cycle through all enabled modes",
    handler: (ctx) => {
      cycleMode(ctx);
    },
  });

  // ─── Flags ─────────────────────────────────────────────────────────────────

  pi.registerFlag("chat", {
    type: "boolean",
    description: "Start in chat mode (read-only conversational)",
  });

  pi.registerFlag("plan", {
    type: "boolean",
    description: "Start in plan mode (read-only, explore and plan)",
  });
}