import type {
  AgentEndEvent,
  ExtensionAPI,
  ExtensionCommandContext,
  ExtensionContext,
  ThinkingLevel,
  ToolCallEvent,
  UserBashEvent,
} from "@mariozechner/pi-coding-agent";
import { DynamicBorder } from "@mariozechner/pi-coding-agent";
import type { Api, Model } from "@mariozechner/pi-ai";
import { Container, Text } from "@mariozechner/pi-tui";
import { existsSync, readFileSync } from "node:fs";
import fs from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import {
  buildAskUserProtocolInstructions,
  formatAskUserResponsePrompt,
  parseAskUserRequestFromText,
  type IAskUserQuestion,
  type IAskUserResponse,
} from "../shared/askUserProtocol";

const PLAN_OUTPUT_ROOT_SEGMENTS = [".docyrus", "plans"] as const;
const PLAN_STATE_TYPE = "plan-session";
const PLAN_ANCHOR_TYPE = "plan-anchor";
const PLAN_CONFIG_FILE = ".pi/plan-policy.json";
const PLAN_WIDGET_KEY = "plan-mode";
const PLAN_BRANCH_LABEL = "plan";
const ALLOWED_TODO_ACTIONS = new Set(["list", "list-all", "get"]);
const ALL_THINKING_LEVELS = new Set(["off", "minimal", "low", "medium", "high", "xhigh"] satisfies ThinkingLevel[]);

const PLAN_MODE_SYSTEM_PROMPT = [
  "You are in Docyrus `/plan` mode.",
  "",
  "This is planning only.",
  "Do not implement code, do not mutate repository files, and do not perform remote Docyrus mutations.",
  "Use inspection, analysis, and verification commands only.",
  "",
  "Respond with an actionable Markdown plan using these exact top-level sections in this order:",
  "## Summary",
  "## Implementation Changes",
  "## Test Plan",
  "## Assumptions",
  "",
  "Keep the plan decision-complete for implementation, call out assumptions explicitly, and do not start implementation.",
].join("\n");

const END_PLAN_SUMMARY_PROMPT = `You are returning from a planning-only branch.

Summarize the plan branch as a concise, implementation-ready plan.

Requirements:
- Do not implement anything.
- Preserve the important decisions made on the plan branch.
- Output Markdown with these exact top-level sections:
  - Summary
  - Implementation Changes
  - Test Plan
  - Assumptions
- Keep it concise and actionable.`;

export interface IPlanModelEntry {
	model: string;
	thinkingLevel?: ThinkingLevel;
}

export interface IPlanProfileOverride {
	match: string;
	models?: IPlanModelEntry[];
}

export interface IPlanPolicy {
	models?: IPlanModelEntry[];
	profiles?: Record<string, IPlanProfileOverride>;
}

export interface IPlanSessionState {
	active: boolean;
	mode?: "plan" | "architect";
	originId?: string;
	branchStartId?: string;
	artifactPath?: string;
	task?: string;
	sourceModel?: string;
	sourceThinking?: ThinkingLevel;
	planModel?: string;
	planThinking?: ThinkingLevel;
	startedAt?: string;
}

export interface IReadPlanPolicyResult {
	policy: IPlanPolicy;
	path?: string;
	error?: string;
}

export type IParseResult<T> = { ok: true; value: T } | { ok: false; error: string };

let currentPlanState: IPlanSessionState | undefined;

export interface IPlanningStartParams {
	pi: ExtensionAPI;
	ctx: ExtensionCommandContext;
	mode: "plan" | "architect";
	task?: string;
	artifactPath: string;
	initialPrompt: string;
	profileName?: string;
}

function expandUserPath(inputPath: string): string {
  if (inputPath === "~") {
    return os.homedir();
  }
  if (inputPath.startsWith("~/")) {
    return path.join(os.homedir(), inputPath.slice(2));
  }
  return inputPath;
}

function resolveAgentRootPath(): string {
  const agentDir = process.env.PI_CODING_AGENT_DIR?.trim();
  return agentDir && agentDir.length > 0 ? expandUserPath(agentDir) : path.join(os.homedir(), ".pi", "agent");
}

function resolveGlobalPlanPolicyPath(): string {
  return path.join(resolveAgentRootPath(), "plan-policy.json");
}

function isRecord(value: unknown): value is Record<string, unknown> {
  return typeof value === "object" && value !== null && !Array.isArray(value);
}

export function parsePlanTask(rawArgs: string): string | null {
  const normalized = rawArgs.trim();
  return normalized.length > 0 ? normalized : null;
}

export function slugifyPlanTask(task: string): string {
  const slug = task
    .toLowerCase()
    .replace(/[^a-z0-9]+/g, "-")
    .replace(/^-+|-+$/g, "")
    .slice(0, 60);

  return slug || "plan";
}

export function formatPlanTimestamp(date: Date): string {
  const year = `${date.getFullYear()}`;
  const month = `${date.getMonth() + 1}`.padStart(2, "0");
  const day = `${date.getDate()}`.padStart(2, "0");
  const hour = `${date.getHours()}`.padStart(2, "0");
  const minute = `${date.getMinutes()}`.padStart(2, "0");
  const second = `${date.getSeconds()}`.padStart(2, "0");
  return `${year}${month}${day}-${hour}${minute}${second}`;
}

export function createPlanArtifactPath(params: {
	cwd: string;
	task: string;
	date?: Date;
}): string {
  const timestamp = formatPlanTimestamp(params.date ?? new Date());
  const slug = slugifyPlanTask(params.task);
  return path.join(params.cwd, ...PLAN_OUTPUT_ROOT_SEGMENTS, `${timestamp}-${slug}.md`);
}

export function parsePlanModelSelector(value: unknown): IParseResult<string> {
  if (typeof value !== "string") {
    return { ok: false, error: "expected model selector provider/modelId" };
  }
  if (value.trim() !== value) {
    return { ok: false, error: "expected model selector provider/modelId" };
  }
  const slashIndex = value.indexOf("/");
  if (slashIndex <= 0 || slashIndex >= value.length - 1) {
    return { ok: false, error: "expected model selector provider/modelId" };
  }
  const provider = value.slice(0, slashIndex);
  const modelId = value.slice(slashIndex + 1);
  if (/\s/.test(provider) || /\s/.test(modelId)) {
    return { ok: false, error: "expected model selector provider/modelId" };
  }
  return { ok: true, value };
}

function parseThinkingLevel(value: unknown): IParseResult<ThinkingLevel> {
  if (typeof value !== "string" || !ALL_THINKING_LEVELS.has(value as ThinkingLevel)) {
    return { ok: false, error: "expected one of: off, minimal, low, medium, high, xhigh" };
  }
  return { ok: true, value: value as ThinkingLevel };
}

function parsePlanModelEntry(value: unknown): IParseResult<IPlanModelEntry> {
  if (typeof value === "string") {
    const parsedSelector = parsePlanModelSelector(value);
    if (!parsedSelector.ok) {
      return parsedSelector;
    }
    return {
      ok: true,
      value: {
        model: parsedSelector.value,
      },
    };
  }

  if (!isRecord(value)) {
    return { ok: false, error: "expected model selector string or { model, thinkingLevel } object" };
  }

  const knownKeys = new Set(["model", "thinkingLevel"]);
  for (const key of Object.keys(value)) {
    if (!knownKeys.has(key)) {
      return { ok: false, error: `unknown key: ${key}` };
    }
  }

  const parsedSelector = parsePlanModelSelector(value.model);
  if (!parsedSelector.ok) {
    return { ok: false, error: `model entry: ${parsedSelector.error}` };
  }

  const entry: IPlanModelEntry = {
    model: parsedSelector.value,
  };

  if (value.thinkingLevel !== undefined) {
    const parsedThinking = parseThinkingLevel(value.thinkingLevel);
    if (!parsedThinking.ok) {
      return { ok: false, error: `model entry thinkingLevel: ${parsedThinking.error}` };
    }
    entry.thinkingLevel = parsedThinking.value;
  }

  return { ok: true, value: entry };
}

function parsePlanModelEntries(value: unknown): IParseResult<IPlanModelEntry[]> {
  if (!Array.isArray(value)) {
    return { ok: false, error: "expected array of model entries" };
  }
  if (value.length === 0) {
    return { ok: false, error: "models array must not be empty" };
  }

  const entries: IPlanModelEntry[] = [];
  for (const item of value) {
    const parsed = parsePlanModelEntry(item);
    if (!parsed.ok) {
      return parsed;
    }
    entries.push(parsed.value);
  }

  return { ok: true, value: entries };
}

function parsePlanProfileOverride(name: string, value: unknown): IParseResult<IPlanProfileOverride> {
  if (!isRecord(value)) {
    return { ok: false, error: `profiles.${name}: expected object` };
  }

  const knownKeys = new Set(["match", "models"]);
  for (const key of Object.keys(value)) {
    if (!knownKeys.has(key)) {
      return { ok: false, error: `profiles.${name}: unknown key: ${key}` };
    }
  }

  const parsedMatch = parsePlanModelSelector(value.match);
  if (!parsedMatch.ok) {
    return { ok: false, error: `profiles.${name}.match: ${parsedMatch.error}` };
  }

  const profile: IPlanProfileOverride = {
    match: parsedMatch.value,
  };

  if (value.models !== undefined) {
    const parsedModels = parsePlanModelEntries(value.models);
    if (!parsedModels.ok) {
      return { ok: false, error: `profiles.${name}.models: ${parsedModels.error}` };
    }
    profile.models = parsedModels.value;
  }

  return { ok: true, value: profile };
}

function parsePlanProfiles(value: unknown): IParseResult<Record<string, IPlanProfileOverride>> {
  if (!isRecord(value)) {
    return { ok: false, error: "profiles: expected object" };
  }

  const profiles: Record<string, IPlanProfileOverride> = {};
  for (const [name, entry] of Object.entries(value)) {
    const parsed = parsePlanProfileOverride(name, entry);
    if (!parsed.ok) {
      return parsed;
    }
    profiles[name] = parsed.value;
  }

  return { ok: true, value: profiles };
}

export function parsePlanPolicy(value: unknown): IParseResult<IPlanPolicy> {
  if (!isRecord(value)) {
    return { ok: false, error: "expected object" };
  }

  const knownKeys = new Set(["models", "profiles"]);
  for (const key of Object.keys(value)) {
    if (!knownKeys.has(key)) {
      return { ok: false, error: `unknown key: ${key}` };
    }
  }

  const policy: IPlanPolicy = {};

  if (value.models !== undefined) {
    const parsedModels = parsePlanModelEntries(value.models);
    if (!parsedModels.ok) {
      return { ok: false, error: `models: ${parsedModels.error}` };
    }
    policy.models = parsedModels.value;
  }

  if (value.profiles !== undefined) {
    const parsedProfiles = parsePlanProfiles(value.profiles);
    if (!parsedProfiles.ok) {
      return parsedProfiles;
    }
    policy.profiles = parsedProfiles.value;
  }

  return { ok: true, value: policy };
}

function readPlanPolicyConfigFile(configPath: string): IReadPlanPolicyResult {
  try {
    const raw = readFileSync(configPath, "utf8");
    const parsed = parsePlanPolicy(JSON.parse(raw) as unknown);
    if (!parsed.ok) {
      return { policy: {}, path: configPath, error: `Invalid ${configPath}: ${parsed.error}` };
    }
    return { policy: parsed.value, path: configPath };
  } catch (error) {
    const message = error instanceof Error ? error.message : String(error);
    return { policy: {}, path: configPath, error: `Invalid ${configPath}: ${message}` };
  }
}

export function readPlanPolicy(cwd: string): IReadPlanPolicyResult {
  const projectPath = path.join(cwd, PLAN_CONFIG_FILE);
  if (existsSync(projectPath)) {
    return readPlanPolicyConfigFile(projectPath);
  }

  const globalPath = resolveGlobalPlanPolicyPath();
  if (existsSync(globalPath)) {
    return readPlanPolicyConfigFile(globalPath);
  }

  return { policy: {} };
}

export function findMatchingPlanProfile(
  profiles: Record<string, IPlanProfileOverride> | undefined,
  sourceModel: string | undefined,
): { name: string; override: IPlanProfileOverride } | undefined {
  if (!profiles || !sourceModel) {
    return undefined;
  }

  for (const name of Object.keys(profiles).sort()) {
    const profile = profiles[name];
    if (profile && profile.match === sourceModel) {
      return { name, override: profile };
    }
  }

  return undefined;
}

export function resolveEffectivePlanPolicy(params: {
	basePolicy: IPlanPolicy;
	sourceModel: string | undefined;
}): {
	profileName: string | undefined;
	modelEntries: IPlanModelEntry[];
} {
  const profile = findMatchingPlanProfile(params.basePolicy.profiles, params.sourceModel);
  return {
    profileName: profile?.name,
    modelEntries: profile?.override.models ?? params.basePolicy.models ?? [],
  };
}

function formatModelSelector(model: { provider: string; id: string }): string {
  return `${model.provider}/${model.id}`;
}

function splitModelSelector(selector: string): IParseResult<{ provider: string; modelId: string }> {
  const parsed = parsePlanModelSelector(selector);
  if (!parsed.ok) {
    return parsed;
  }

  const slashIndex = parsed.value.indexOf("/");
  return {
    ok: true,
    value: {
      provider: parsed.value.slice(0, slashIndex),
      modelId: parsed.value.slice(slashIndex + 1),
    },
  };
}

function getCurrentModelSelector(ctx: ExtensionContext): string | undefined {
  return ctx.model ? formatModelSelector(ctx.model) : undefined;
}

async function tryResolveConfiguredModel(
  ctx: ExtensionContext,
  selector: string,
): Promise<Model<Api> | undefined> {
  const parsed = splitModelSelector(selector);
  if (!parsed.ok) {
    return undefined;
  }

  const model = ctx.modelRegistry.find(parsed.value.provider, parsed.value.modelId);
  if (!model) {
    return undefined;
  }

  try {
    const apiKey = await ctx.modelRegistry.getApiKey(model);
    if (!apiKey) {
      return undefined;
    }
  } catch {
    return undefined;
  }

  return model as Model<Api>;
}

async function resolveConfiguredPlanModel(
  ctx: ExtensionContext,
  modelEntries: IPlanModelEntry[],
): Promise<{ entry: IPlanModelEntry; model: Model<Api> } | undefined> {
  for (const entry of modelEntries) {
    const model = await tryResolveConfiguredModel(ctx, entry.model);
    if (model) {
      return { entry, model };
    }
  }

  return undefined;
}

function buildPlanPromptOverlay(state: IPlanSessionState): string {
  const task = state.task?.trim() || "(unknown task)";
  const model = state.planModel ? `Planning model: ${state.planModel}` : "Planning model: current session model";
  const artifact = state.artifactPath ? `Artifact path: ${state.artifactPath}` : "Artifact path: unavailable";
  const modeLabel = state.mode === "architect" ? "Docyrus `/architect` planning mode." : "Docyrus `/plan` mode.";
  return [
    modeLabel,
    PLAN_MODE_SYSTEM_PROMPT,
    "",
    buildAskUserProtocolInstructions(),
    "",
    "Current planning session:",
    `Task: ${task}`,
    model,
    artifact,
  ].join("\n");
}

function buildInitialPlanPrompt(task: string | undefined): string {
  if (!task) {
    return [
      "The user started a planning workflow but did not provide a concrete task yet.",
      "",
      "Before drafting any plan, gather the minimum clarification you need from the user.",
      "If clarification is needed, use the ask_user protocol.",
    ].join("\n");
  }

  return [
    "Create a decision-complete implementation plan for the following task.",
    "",
    task,
  ].join("\n");
}

function buildInitialArtifactContent(task: string): string {
  return [
    "# Planning Session",
    "",
    "## Task",
    task,
    "",
    "_Waiting for the first planning response..._",
    "",
  ].join("\n");
}

function getPlanState(ctx: ExtensionContext): IPlanSessionState | undefined {
  let state: IPlanSessionState | undefined;
  for (const entry of ctx.sessionManager.getBranch()) {
    if (entry.type === "custom" && entry.customType === PLAN_STATE_TYPE) {
      state = entry.data as IPlanSessionState | undefined;
    }
  }

  if (state?.active) {
    return state;
  }

  return undefined;
}

function setPlanWidget(ctx: ExtensionContext, state: IPlanSessionState | undefined): void {
  if (!ctx.hasUI) {
    return;
  }

  if (!state?.active) {
    ctx.ui.setWidget(PLAN_WIDGET_KEY, undefined);
    return;
  }

  ctx.ui.setWidget(PLAN_WIDGET_KEY, (_tui, theme) => {
    const lines = [
      theme.fg("accent", theme.bold("Plan mode active")),
      theme.fg("muted", `Model: ${state.planModel ?? "current session model"}`),
      theme.fg("muted", `Artifact: ${state.artifactPath ?? "(none)"}`),
      theme.fg("warning", "Exit with /end-plan"),
    ];

    const container = new Container();
    container.addChild(new DynamicBorder((value: string) => theme.fg("accent", value)));
    container.addChild(new Text(lines.join("\n"), 0, 0));
    return container;
  });
}

async function ensureArtifactFile(artifactPath: string, task: string): Promise<void> {
  await fs.mkdir(path.dirname(artifactPath), { recursive: true });
  await fs.writeFile(artifactPath, buildInitialArtifactContent(task), "utf8");
}

function appendPlanState(pi: ExtensionAPI, state: IPlanSessionState): void {
  pi.appendEntry(PLAN_STATE_TYPE, state);
}

function clearPlanState(pi: ExtensionAPI): void {
  pi.appendEntry(PLAN_STATE_TYPE, { active: false });
}

function getLatestUserEntryId(ctx: ExtensionContext): string | undefined {
  const branch = ctx.sessionManager.getBranch();
  for (let index = branch.length - 1; index >= 0; index -= 1) {
    const entry = branch[index];
    if (entry.type === "message" && entry.message.role === "user") {
      return entry.id;
    }
  }

  return undefined;
}

function extractLastAssistantText(messages: Array<{ role?: string; content?: unknown }>): string | null {
  for (let index = messages.length - 1; index >= 0; index -= 1) {
    const message = messages[index];
    if (message?.role !== "assistant") {
      continue;
    }

    if (typeof message.content === "string") {
      return message.content.trim() || null;
    }

    if (Array.isArray(message.content)) {
      const text = message.content
        .filter((block): block is { type: "text"; text: string } => {
          return Boolean(
            block &&
						typeof block === "object" &&
						"type" in block &&
						"text" in block &&
						block.type === "text" &&
						typeof block.text === "string",
          );
        })
        .map((block) => block.text)
        .join("\n")
        .trim();

      return text || null;
    }

    return null;
  }

  return null;
}

async function writePlanArtifactFromEvent(event: AgentEndEvent, ctx: ExtensionContext): Promise<void> {
  const state = getPlanState(ctx);
  if (!state?.active || !state.artifactPath || state.mode === "architect") {
    return;
  }

  const text = extractLastAssistantText(event.messages ?? []);
  if (!text || parseAskUserRequestFromText(text)) {
    return;
  }

  try {
    await fs.mkdir(path.dirname(state.artifactPath), { recursive: true });
    await fs.writeFile(state.artifactPath, `${text.trim()}\n`, "utf8");
  } catch (error) {
    if (ctx.hasUI) {
      const message = error instanceof Error ? error.message : String(error);
      ctx.ui.notify(`Failed to write plan artifact: ${message}`, "error");
    }
  }
}

export function shouldBlockTodoAction(action: string | undefined): boolean {
  return !!action && !ALLOWED_TODO_ACTIONS.has(action);
}

export function isPotentiallyMutatingShellCommand(command: string): boolean {
  const trimmed = command.trim();
  if (!trimmed) {
    return false;
  }

  const patterns = [
    /(^|[;&|]\s*)git\s+(add|commit|push|pull|merge|rebase|cherry-pick|reset|clean|restore|checkout|switch|stash|am|apply)\b/i,
    /(^|[;&|]\s*)(rm|mv|cp|mkdir|touch|install|chmod|chown)\b/i,
    /(^|[;&|]\s*)(npm|pnpm|yarn|bun)\s+(add|install|remove|unlink|link|update|upgrade|publish|version)\b/i,
    /(^|[;&|]\s*)(pip|pip3)\s+(install|uninstall)\b/i,
    /(^|[;&|]\s*)cargo\s+(add|install|remove)\b/i,
    /(^|[;&|]\s*)go\s+(get|mod\s+tidy)\b/i,
    /(^|[;&|]\s*)brew\s+(install|upgrade|tap|untap)\b/i,
    /(^|[;&|]\s*)sed\s+-i\b/i,
    /(^|[;&|]\s*)perl\s+-i\b/i,
    /(^|[;&|]\s*)tee\b/i,
    /(^|[;&|]\s*)docyrus\s+(apps|ds|studio)\s+(create|update|delete|restore|permanent-delete|bulk-create-data-sources|create-field|update-field|delete-field|create-fields-batch|update-fields-batch|delete-fields-batch|create-enums|update-enums|delete-enums)\b/i,
  ];

  if (patterns.some((pattern) => pattern.test(trimmed))) {
    return true;
  }

  if (/\s(?:\d*>>?)(?:\s|$)/.test(trimmed)) {
    return true;
  }

  return false;
}

function buildBlockedToolReason(detail: string): string {
  return `A planning session is active. ${detail} Exit with /end-plan or /end-architect before implementing.`;
}

function buildBlockedUserBashResult(detail: string) {
  return {
    result: {
      output: `${buildBlockedToolReason(detail)}\n`,
      exitCode: 1,
      cancelled: false,
      truncated: false,
    },
  };
}

async function applyModelSelector(
  pi: ExtensionAPI,
  ctx: ExtensionContext,
  selector: string | undefined,
  thinkingLevel: ThinkingLevel | undefined,
): Promise<boolean> {
  if (!selector) {
    if (thinkingLevel) {
      pi.setThinkingLevel(thinkingLevel);
    }
    return true;
  }

  const parts = splitModelSelector(selector);
  if (!parts.ok) {
    if (ctx.hasUI) {
      ctx.ui.notify(`Invalid plan model selector ${selector}`, "warning");
    }
    return false;
  }

  const model = ctx.modelRegistry.find(parts.value.provider, parts.value.modelId);
  if (!model) {
    if (ctx.hasUI) {
      ctx.ui.notify(`Plan model ${selector} is not available in the current registry`, "warning");
    }
    return false;
  }

  const changed = getCurrentModelSelector(ctx) !== selector;
  const ok = changed ? await pi.setModel(model) : true;
  if (!ok) {
    if (ctx.hasUI) {
      ctx.ui.notify(`No API key available for plan model ${selector}`, "warning");
    }
    return false;
  }

  if (thinkingLevel) {
    pi.setThinkingLevel(thinkingLevel);
  }

  return true;
}

async function applyResolvedPlanModel(
  pi: ExtensionAPI,
  ctx: ExtensionContext,
  resolved: { entry: IPlanModelEntry; model: Model<Api> } | undefined,
): Promise<{ planModel?: string; planThinking?: ThinkingLevel }> {
  if (!resolved) {
    return {
      planModel: getCurrentModelSelector(ctx),
      planThinking: pi.getThinkingLevel(),
    };
  }

  const ok = await pi.setModel(resolved.model);
  if (!ok) {
    if (ctx.hasUI) {
      ctx.ui.notify(`No API key available for plan model ${resolved.entry.model}`, "warning");
    }
    return {
      planModel: getCurrentModelSelector(ctx),
      planThinking: pi.getThinkingLevel(),
    };
  }

  if (resolved.entry.thinkingLevel) {
    pi.setThinkingLevel(resolved.entry.thinkingLevel);
  }

  return {
    planModel: formatModelSelector(resolved.model),
    planThinking: resolved.entry.thinkingLevel ?? pi.getThinkingLevel(),
  };
}

async function restoreSourceModel(
  pi: ExtensionAPI,
  ctx: ExtensionContext,
  state: IPlanSessionState | undefined,
): Promise<void> {
  if (!state) {
    return;
  }

  const restored = await applyModelSelector(pi, ctx, state.sourceModel, state.sourceThinking);
  if (!restored && state.sourceThinking) {
    pi.setThinkingLevel(state.sourceThinking);
  }
}

async function syncPlanState(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
  const nextState = getPlanState(ctx);
  const previousState = currentPlanState;
  currentPlanState = nextState;
  setPlanWidget(ctx, nextState);

  if (nextState?.active) {
    const currentSelector = getCurrentModelSelector(ctx);
    const currentThinking = pi.getThinkingLevel();
    if (nextState.planModel !== currentSelector || nextState.planThinking !== currentThinking) {
      await applyModelSelector(pi, ctx, nextState.planModel, nextState.planThinking);
    }
    return;
  }

  if (previousState?.active) {
    await restoreSourceModel(pi, ctx, previousState);
  }
}

function formatModelChain(entries: IPlanModelEntry[]): string {
  if (entries.length === 0) {
    return "(none)";
  }
  return entries
    .map((entry) => entry.thinkingLevel ? `${entry.model} (thinking: ${entry.thinkingLevel})` : entry.model)
    .join(" -> ");
}

export async function startPlanningWorkflow(params: IPlanningStartParams): Promise<boolean> {
  const { pi, ctx, mode, artifactPath, initialPrompt } = params;
  if (getPlanState(ctx)) {
    if (ctx.hasUI) {
      ctx.ui.notify("A planning session is already active. End it before starting another one.", "warning");
    }
    return false;
  }

  const sourceModel = getCurrentModelSelector(ctx);
  const sourceThinking = pi.getThinkingLevel();
  const policyResult = readPlanPolicy(ctx.cwd);
  if (policyResult.error && ctx.hasUI) {
    ctx.ui.notify(policyResult.error, "warning");
  }

  const { profileName, modelEntries } = resolveEffectivePlanPolicy({
    basePolicy: policyResult.policy,
    sourceModel,
  });
  const resolvedPlanModel = await resolveConfiguredPlanModel(ctx, modelEntries);
  if (!resolvedPlanModel && modelEntries.length > 0) {
    ctx.ui.notify(
      `No configured planning models could be resolved (${formatModelChain(modelEntries)}). Staying on the current model.`,
      "warning",
    );
  }

  try {
    await ensureArtifactFile(artifactPath, params.task ?? mode);
  } catch (error) {
    if (ctx.hasUI) {
      const message = error instanceof Error ? error.message : String(error);
      ctx.ui.notify(`Failed to create plan artifact ${artifactPath}: ${message}`, "error");
    }
    return false;
  }

  let originId = ctx.sessionManager.getLeafId() ?? undefined;
  if (!originId) {
    pi.appendEntry(PLAN_ANCHOR_TYPE, { createdAt: new Date().toISOString() });
    originId = ctx.sessionManager.getLeafId() ?? undefined;
  }

  if (!originId) {
    if (ctx.hasUI) {
      ctx.ui.notify("Failed to determine the current branch position for the planning workflow.", "error");
    }
    return false;
  }

  let branchStartTargetId = getLatestUserEntryId(ctx);
  if (branchStartTargetId && branchStartTargetId === originId) {
    pi.appendEntry(PLAN_ANCHOR_TYPE, { createdAt: new Date().toISOString() });
    originId = ctx.sessionManager.getLeafId() ?? originId;
  }

  if (branchStartTargetId && branchStartTargetId !== originId) {
    try {
      const result = await ctx.navigateTree(branchStartTargetId, {
        summarize: false,
        label: PLAN_BRANCH_LABEL,
      });
      if (result.cancelled) {
        return false;
      }
    } catch (error) {
      if (ctx.hasUI) {
        const message = error instanceof Error ? error.message : String(error);
        ctx.ui.notify(`Failed to start the planning branch: ${message}`, "error");
      }
      return false;
    }
  }

  const branchStartId = ctx.sessionManager.getLeafId() ?? branchStartTargetId;
  const appliedPlanModel = await applyResolvedPlanModel(pi, ctx, resolvedPlanModel);
  const state: IPlanSessionState = {
    active: true,
    mode,
    originId,
    branchStartId,
    artifactPath,
    task: params.task,
    sourceModel,
    sourceThinking,
    planModel: appliedPlanModel.planModel,
    planThinking: appliedPlanModel.planThinking,
    startedAt: new Date().toISOString(),
  };

  appendPlanState(pi, state);
  currentPlanState = state;
  setPlanWidget(ctx, state);

  if (ctx.hasUI) {
    const profileHint = profileName ? ` (profile: ${profileName})` : "";
    const modeLabel = mode === "architect" ? "Architect mode" : "Plan mode";
    ctx.ui.notify(`${modeLabel} active${profileHint}. Writing updates to ${artifactPath}`, "info");
  }

  pi.sendUserMessage(initialPrompt);
  return true;
}

async function planHandler(pi: ExtensionAPI, ctx: ExtensionCommandContext, args: string): Promise<void> {
  const task = parsePlanTask(args) ?? undefined;
  const artifactPath = createPlanArtifactPath({
    cwd: ctx.cwd,
    task: task ?? "plan",
    date: new Date(),
  });

  await startPlanningWorkflow({
    pi,
    ctx,
    mode: "plan",
    task,
    artifactPath,
    initialPrompt: buildInitialPlanPrompt(task),
  });
}

export async function endPlanningWorkflow(pi: ExtensionAPI, ctx: ExtensionCommandContext): Promise<void> {
  const state = getPlanState(ctx);
  if (!state?.active || !state.originId) {
    if (ctx.hasUI) {
      ctx.ui.notify("No active planning session was found on the current branch.", "info");
    }
    return;
  }

  clearPlanState(pi);

  try {
    const result = await ctx.navigateTree(state.originId, {
      summarize: true,
      customInstructions: END_PLAN_SUMMARY_PROMPT,
      replaceInstructions: true,
      label: PLAN_BRANCH_LABEL,
    });

    if (result.cancelled) {
      currentPlanState = undefined;
      setPlanWidget(ctx, undefined);
      await restoreSourceModel(pi, ctx, state);
      if (ctx.hasUI) {
        ctx.ui.notify("Planning session ended in place after navigation was cancelled.", "info");
      }
      return;
    }
  } catch (error) {
    currentPlanState = undefined;
    setPlanWidget(ctx, undefined);
    await restoreSourceModel(pi, ctx, state);
    if (ctx.hasUI) {
      const message = error instanceof Error ? error.message : String(error);
      ctx.ui.notify(`Failed to return from the planning branch: ${message}`, "error");
    }
    return;
  }

  currentPlanState = undefined;
  setPlanWidget(ctx, undefined);
  await restoreSourceModel(pi, ctx, state);

  if (ctx.hasUI) {
    const artifactSuffix = state.artifactPath ? ` Artifact: ${state.artifactPath}` : "";
    ctx.ui.notify(`Planning session ended.${artifactSuffix}`, "info");
  }
}

async function collectAskUserResponse(ctx: ExtensionContext, request: { title: string; message: string; questions: IAskUserQuestion[] }): Promise<IAskUserResponse | undefined> {
  if (!ctx.hasUI) {
    return undefined;
  }

  const answers: Record<string, string | boolean | string[]> = {};

  for (const question of request.questions) {
    if (question.type === "boolean") {
      const confirmed = await ctx.ui.confirm(question.label, question.description ?? request.message);
      answers[question.id] = confirmed;
      continue;
    }

    if (question.type === "singleSelect" && question.options && question.options.length > 0) {
      const selectedLabel = await ctx.ui.select(
        question.label,
        question.options.map((option) => option.label),
      );
      const selected = question.options.find((option) => option.label === selectedLabel);
      if (selected) {
        answers[question.id] = selected.value;
      }
      continue;
    }

    const prefill = Array.isArray(question.defaultValue)
      ? question.defaultValue.join(", ")
      : typeof question.defaultValue === "string"
        ? question.defaultValue
        : "";
    const input = await ctx.ui.editor(question.label, prefill);
    if (input === undefined) {
      return {
        action: "cancel",
        answers,
      };
    }

    if (question.type === "multiSelect") {
      answers[question.id] = input
        .split(",")
        .map((value) => value.trim())
        .filter(Boolean);
    } else {
      answers[question.id] = input.trim();
    }
  }

  return {
    action: "submit",
    answers,
  };
}

async function planPolicyHandler(_pi: ExtensionAPI, ctx: ExtensionCommandContext): Promise<void> {
  const branchState = getPlanState(ctx);
  const sourceModel = branchState?.sourceModel ?? getCurrentModelSelector(ctx);
  const currentModel = getCurrentModelSelector(ctx);
  const policyResult = readPlanPolicy(ctx.cwd);
  if (policyResult.error) {
    ctx.ui.notify(policyResult.error, "warning");
  }

  const { profileName, modelEntries } = resolveEffectivePlanPolicy({
    basePolicy: policyResult.policy,
    sourceModel,
  });
  const resolved = await resolveConfiguredPlanModel(ctx, modelEntries);
  const resolvedSelector = resolved ? formatModelSelector(resolved.model) : branchState?.planModel ?? currentModel ?? "(unresolved)";
  const resolvedThinking = resolved?.entry.thinkingLevel ?? branchState?.planThinking;

  const lines = [
    `status: ${branchState?.active ? "active" : "inactive"}`,
    `config: ${policyResult.path ?? "(none)"}`,
    `source model: ${sourceModel ?? "unknown"}`,
    `current model: ${currentModel ?? "unknown"}`,
    `profile: ${profileName ?? "none"}`,
    `configured models: ${formatModelChain(modelEntries)}`,
    `resolved planning model: ${resolvedThinking ? `${resolvedSelector} (thinking: ${resolvedThinking})` : resolvedSelector}`,
  ];

  if (branchState?.artifactPath) {
    lines.push(`artifact: ${branchState.artifactPath}`);
  }

  ctx.ui.notify(lines.join("\n"), "info");
}

function handleToolCallDuringPlan(event: ToolCallEvent, ctx: ExtensionContext): { block: boolean; reason: string } | undefined {
  if (!getPlanState(ctx)?.active) {
    return undefined;
  }

  if (event.toolName === "edit" || event.toolName === "write") {
    return {
      block: true,
      reason: buildBlockedToolReason(`The ${event.toolName} tool is disabled during planning.`),
    };
  }

  if (event.toolName === "bash" && isRecord(event.input) && typeof event.input.command === "string") {
    if (isPotentiallyMutatingShellCommand(event.input.command)) {
      return {
        block: true,
        reason: buildBlockedToolReason(`Mutating bash commands are disabled during planning (${event.input.command}).`),
      };
    }
  }

  if (event.toolName === "todo" && isRecord(event.input) && typeof event.input.action === "string") {
    if (shouldBlockTodoAction(event.input.action)) {
      return {
        block: true,
        reason: buildBlockedToolReason(`Todo action "${event.input.action}" is disabled during planning.`),
      };
    }
  }

  return undefined;
}

function handleUserBashDuringPlan(event: UserBashEvent, ctx: ExtensionContext) {
  if (!getPlanState(ctx)?.active) {
    return undefined;
  }

  if (!isPotentiallyMutatingShellCommand(event.command)) {
    return undefined;
  }

  return buildBlockedUserBashResult(`Mutating shell commands are disabled during planning (${event.command}).`);
}

export default function planExtension(pi: ExtensionAPI) {
  pi.registerCommand("plan", {
    description: "Start a planning-only branch with an optional dedicated planning model",
    handler: async(args, ctx) => {
      await planHandler(pi, ctx, args);
    },
  });

  pi.registerCommand("end-plan", {
    description: "Leave plan mode, summarize the plan branch, and return to the original branch",
    handler: async(_args, ctx) => {
      await endPlanningWorkflow(pi, ctx);
    },
  });

  pi.registerCommand("end-architect", {
    description: "Leave architect mode, summarize the architect branch, and return to the original branch",
    handler: async(_args, ctx) => {
      await endPlanningWorkflow(pi, ctx);
    },
  });

  pi.registerCommand("plan-policy", {
    description: "Show effective planning-model policy and the resolved planning model",
    handler: async(_args, ctx) => {
      await planPolicyHandler(pi, ctx);
    },
  });

  pi.on("before_agent_start", (event, ctx) => {
    const state = getPlanState(ctx);
    if (!state?.active) {
      return;
    }

    return {
      systemPrompt: [event.systemPrompt, buildPlanPromptOverlay(state)].filter(Boolean).join("\n\n"),
    };
  });

  pi.on("tool_call", (event, ctx) => {
    return handleToolCallDuringPlan(event, ctx);
  });

  pi.on("user_bash", (event, ctx) => {
    return handleUserBashDuringPlan(event, ctx);
  });

  pi.on("agent_end", async(event, ctx) => {
    const state = getPlanState(ctx);
    if (state?.active) {
      const text = extractLastAssistantText(event.messages ?? []);
      const askUserRequest = text ? parseAskUserRequestFromText(text) : undefined;
      if (askUserRequest && ctx.hasUI) {
        const response = await collectAskUserResponse(ctx, askUserRequest);
        if (response) {
          pi.sendUserMessage(formatAskUserResponsePrompt(response));
        }
        return;
      }
    }

    await writePlanArtifactFromEvent(event, ctx);
  });

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

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

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

  pi.on("session_shutdown", async(_event, ctx) => {
    currentPlanState = undefined;
    setPlanWidget(ctx, undefined);
  });
}
