/**
 * Plan Mode Extension
 *
 * Read-only exploration + structured planning + gated step-by-step execution.
 *
 * Features:
 * - /plan command or Ctrl+Alt+P to toggle read-only plan mode
 * - Bash restricted to allowlisted read-only commands, edit/write disabled
 * - `submit_plan` tool: model returns a STRUCTURED plan (title + steps + risks + files)
 * - Plans are saved to `.pi/plans/<slug>-<timestamp>.md` (recognizable name + timestamp)
 * - Per-step execution gating: you approve/skip/stop each step before it runs
 * - Progress tracking widget + footer status, persisted across resume
 */

import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { TextContent } from "@earendil-works/pi-ai";
import {
	CONFIG_DIR_NAME,
	type ExtensionAPI,
	type ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Key } from "@earendil-works/pi-tui";
import { mkdir, writeFile } from "node:fs/promises";
import { join } from "node:path";
import { Type, type Static } from "typebox";
import { isSafeCommand } from "./utils.ts";

// Tools
const PLAN_MODE_TOOLS = ["read", "bash", "grep", "find", "ls", "questionnaire", "submit_plan"];
const NORMAL_MODE_TOOLS = ["read", "bash", "edit", "write"];
const PLAN_MODE_DISABLED_TOOLS = new Set<string>(["edit", "write"]);
const PLAN_MANAGED_TOOLS = new Set<string>([...PLAN_MODE_TOOLS, ...NORMAL_MODE_TOOLS]);

type Phase = "planning" | "awaiting-approval" | "executing";

interface PlanStep {
	text: string;
	completed?: boolean;
	skipped?: boolean;
}

interface StoredPlan {
	title: string;
	createdAt: string;
	steps: PlanStep[];
	risks?: string[];
	files?: string[];
}

interface PlanModeState {
	enabled: boolean;
	phase?: Phase;
	plan?: StoredPlan;
	currentStepIndex?: number;
	planFilePath?: string;
	toolsBeforePlanMode?: string[];
}

// ---- submit_plan tool schema ----
const submitPlanSchema = Type.Object({
	title: Type.String({
		description: "A short, recognizable name for this plan (e.g. 'Add OAuth login').",
	}),
	steps: Type.Array(
		Type.Object({
			description: Type.String({ description: "A single, concrete, actionable step." }),
		}),
		{ description: "Ordered list of steps to implement the plan." },
	),
	risks: Type.Optional(Type.Array(Type.String(), { description: "Potential risks, caveats, or open questions." })),
	files: Type.Optional(Type.Array(Type.String(), { description: "Files likely to be created or modified." })),
});
export type SubmitPlanInput = Static<typeof submitPlanSchema>;

// ---- file persistence helpers ----
function slugify(title: string): string {
	const slug = title
		.toLowerCase()
		.replace(/[^a-z0-9]+/g, "-")
		.replace(/^-+|-+$/g, "")
		.slice(0, 50)
		.replace(/-+$/g, "");
	return slug || "plan";
}

function fileTimestamp(): string {
	const d = new Date();
	const p = (n: number) => String(n).padStart(2, "0");
	return `${d.getFullYear()}${p(d.getMonth() + 1)}${p(d.getDate())}-${p(d.getHours())}${p(d.getMinutes())}${p(d.getSeconds())}`;
}

function renderPlanMarkdown(plan: StoredPlan): string {
	const lines: string[] = [];
	lines.push(`# ${plan.title}`);
	lines.push("");
	lines.push(`_Created: ${plan.createdAt}_`);
	lines.push("");
	lines.push("## Steps");
	lines.push("");
	plan.steps.forEach((s, i) => {
		const box = s.completed ? "x" : " ";
		const suffix = s.skipped ? " _(skipped)_" : "";
		lines.push(`${i + 1}. [${box}] ${s.text}${suffix}`);
	});
	if (plan.risks?.length) {
		lines.push("");
		lines.push("## Risks");
		lines.push("");
		for (const r of plan.risks) lines.push(`- ${r}`);
	}
	if (plan.files?.length) {
		lines.push("");
		lines.push("## Files");
		lines.push("");
		for (const f of plan.files) lines.push(`- ${f}`);
	}
	lines.push("");
	return lines.join("\n");
}

export default function planModeExtension(pi: ExtensionAPI): void {
	let planModeEnabled = false;
	let phase: Phase | undefined;
	let plan: StoredPlan | undefined;
	let currentStepIndex = 0;
	let awaitingStep = false;
	let planFilePath: string | undefined;
	let toolsBeforePlanMode: string[] | undefined;

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

	// ---- tool management ----
	function uniqueToolNames(toolNames: string[]): string[] {
		return [...new Set(toolNames)];
	}
	function getPlanModeTools(activeToolNames: string[]): string[] {
		return uniqueToolNames([
			...activeToolNames.filter((name) => !PLAN_MODE_DISABLED_TOOLS.has(name)),
			...PLAN_MODE_TOOLS,
		]);
	}
	function getNormalModeTools(activeToolNames: string[]): string[] {
		return uniqueToolNames([
			...NORMAL_MODE_TOOLS,
			...activeToolNames.filter((name) => !PLAN_MANAGED_TOOLS.has(name)),
		]);
	}
	function enablePlanModeTools(): void {
		if (toolsBeforePlanMode === undefined) {
			toolsBeforePlanMode = pi.getActiveTools();
		}
		pi.setActiveTools(getPlanModeTools(toolsBeforePlanMode));
	}
	function restoreNormalModeTools(): void {
		pi.setActiveTools(toolsBeforePlanMode ?? getNormalModeTools(pi.getActiveTools()));
		toolsBeforePlanMode = undefined;
	}

	// ---- state ----
	function persistState(): void {
		pi.appendEntry("plan-mode", {
			enabled: planModeEnabled,
			phase,
			plan,
			currentStepIndex,
			planFilePath,
			toolsBeforePlanMode,
		} satisfies PlanModeState);
	}

	async function persistPlanFile(): Promise<void> {
		if (!planFilePath || !plan) return;
		try {
			await writeFile(planFilePath, renderPlanMarkdown(plan), "utf8");
		} catch {
			// ignore write failures
		}
	}

	async function savePlanFile(ctx: ExtensionContext, p: StoredPlan): Promise<string> {
		const dir = join(ctx.cwd, CONFIG_DIR_NAME, "plans");
		await mkdir(dir, { recursive: true });
		const path = join(dir, `${slugify(p.title)}-${fileTimestamp()}.md`);
		await writeFile(path, renderPlanMarkdown(p), "utf8");
		return path;
	}

	// ---- UI ----
	function updateStatus(ctx: ExtensionContext): void {
		if (phase === "executing" && plan) {
			const done = plan.steps.filter((s) => s.completed).length;
			ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("accent", `📋 ${done}/${plan.steps.length}`));
		} else if (planModeEnabled) {
			ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("warning", "⏸ plan"));
		} else {
			ctx.ui.setStatus("plan-mode", undefined);
		}

		if ((phase === "executing" || phase === "awaiting-approval") && plan) {
			const lines = [ctx.ui.theme.fg("accent", plan.title)];
			plan.steps.forEach((s, i) => {
				let mark: string;
				let text = s.text;
				if (s.completed) {
					mark = ctx.ui.theme.fg("success", "☑ ");
					text = ctx.ui.theme.fg("muted", ctx.ui.theme.strikethrough(text));
				} else if (s.skipped) {
					mark = ctx.ui.theme.fg("muted", "⊘ ");
					text = ctx.ui.theme.fg("muted", text);
				} else if (phase === "executing" && i === currentStepIndex) {
					mark = ctx.ui.theme.fg("accent", "▶ ");
				} else {
					mark = ctx.ui.theme.fg("muted", "☐ ");
				}
				lines.push(`${mark}${text}`);
			});
			ctx.ui.setWidget("plan-todos", lines);
		} else {
			ctx.ui.setWidget("plan-todos", undefined);
		}
	}

	// ---- mode toggle ----
	function togglePlanMode(ctx: ExtensionContext): void {
		planModeEnabled = !planModeEnabled;
		phase = planModeEnabled ? "planning" : undefined;
		plan = undefined;
		planFilePath = undefined;
		currentStepIndex = 0;
		awaitingStep = false;

		if (planModeEnabled) {
			enablePlanModeTools();
			ctx.ui.notify("Plan mode enabled. Built-in write tools disabled. Ask for a plan, then submit_plan.");
		} else {
			restoreNormalModeTools();
			ctx.ui.notify("Plan mode disabled. Full access restored.");
		}
		updateStatus(ctx);
		persistState();
	}

	// ---- execution gating ----
	function finishExecution(ctx: ExtensionContext): void {
		if (plan) {
			const summary = plan.steps
				.map((s) => (s.completed ? `~~${s.text}~~` : s.skipped ? `${s.text} (skipped)` : s.text))
				.join("\n");
			pi.sendMessage(
				{ customType: "plan-complete", content: `**Plan "${plan.title}" complete!** ✓\n\n${summary}`, display: true },
				{ triggerTurn: false },
			);
		}
		phase = undefined;
		awaitingStep = false;
		updateStatus(ctx);
		persistState();
		void persistPlanFile();
	}

	async function gateNext(ctx: ExtensionContext): Promise<void> {
		if (!plan) {
			phase = undefined;
			return;
		}
		if (currentStepIndex >= plan.steps.length) {
			finishExecution(ctx);
			return;
		}

		const step = plan.steps[currentStepIndex];
		if (!step) {
			finishExecution(ctx);
			return;
		}

		const choice = await ctx.ui.select(
			`Step ${currentStepIndex + 1}/${plan.steps.length}: ${step.text}`,
			["Execute this step", "Skip this step", "Stop execution"],
		);

		if (choice?.startsWith("Execute")) {
			awaitingStep = true;
			updateStatus(ctx);
			const stepMsg = `Execute ONLY step ${currentStepIndex + 1} of ${plan.steps.length} from plan "${plan.title}":

${step.text}

Do just this one step. Do NOT start any other step. Stop when this step is done.`;
			pi.sendMessage(
				{ customType: "plan-step", content: stepMsg, display: true },
				{ triggerTurn: true, deliverAs: "followUp" },
			);
		} else if (choice?.startsWith("Skip")) {
			step.skipped = true;
			currentStepIndex++;
			updateStatus(ctx);
			persistState();
			await persistPlanFile();
			await gateNext(ctx);
		} else {
			ctx.ui.notify("Plan execution stopped.", "info");
			phase = undefined;
			awaitingStep = false;
			updateStatus(ctx);
			persistState();
		}
	}

	async function handleApproval(ctx: ExtensionContext): Promise<void> {
		if (!plan) {
			phase = undefined;
			return;
		}
		const choice = await ctx.ui.select(
			`Plan "${plan.title}" — ${plan.steps.length} steps.\nWhat next?`,
			["Execute step-by-step", "Stay in plan mode", "Refine the plan", "Save plan to file"],
		);

		if (choice?.startsWith("Execute")) {
			planModeEnabled = false;
			phase = "executing";
			currentStepIndex = 0;
			awaitingStep = false;
			restoreNormalModeTools();
			updateStatus(ctx);
			persistState();
			await gateNext(ctx);
		} else if (choice === "Refine the plan") {
			phase = "planning";
			const refinement = await ctx.ui.editor("Refine the plan:", "");
			updateStatus(ctx);
			persistState();
			if (refinement?.trim()) {
				pi.sendUserMessage(refinement.trim(), { deliverAs: "followUp" });
			}
		} else if (choice === "Save plan to file") {
			try {
				planFilePath = await savePlanFile(ctx, plan);
				ctx.ui.notify(`Plan saved to ${planFilePath}`, "info");
				persistState();
			} catch {
				ctx.ui.notify("Failed to save plan.", "error");
			}
		} else {
			phase = "planning";
			updateStatus(ctx);
			persistState();
		}
	}

	// ---- commands & shortcuts ----
	pi.registerCommand("plan", {
		description: "Toggle plan mode (read-only exploration + structured planning)",
		handler: async (_args, ctx) => togglePlanMode(ctx),
	});

	pi.registerCommand("plan-status", {
		description: "Show the current plan and progress",
		handler: async (_args, ctx) => {
			if (!plan) {
				ctx.ui.notify("No plan yet. Enable /plan and ask the agent to submit_plan.", "info");
				return;
			}
			const list = plan.steps
				.map((s, i) => `${i + 1}. ${s.completed ? "✓" : s.skipped ? "⊘" : "○"} ${s.text}`)
				.join("\n");
			const loc = planFilePath ? `\n\nFile: ${planFilePath}` : "";
			ctx.ui.notify(`Plan "${plan.title}" (${phase ?? "idle"}):\n${list}${loc}`, "info");
		},
	});

	pi.registerCommand("plan-continue", {
		description: "Resume gated step-by-step execution of the current plan",
		handler: async (_args, ctx) => {
			if (phase !== "executing" || !plan) {
				ctx.ui.notify("No plan execution in progress.", "info");
				return;
			}
			if (awaitingStep) {
				ctx.ui.notify("A step is already running.", "info");
				return;
			}
			await gateNext(ctx);
		},
	});

	pi.registerCommand("plan-save", {
		description: "Save the current plan to a file in .pi/plans/",
		handler: async (_args, ctx) => {
			if (!plan) {
				ctx.ui.notify("No plan to save. Submit a plan first.", "info");
				return;
			}
			try {
				planFilePath = await savePlanFile(ctx, plan);
				ctx.ui.notify(`Plan saved to ${planFilePath}`, "info");
				persistState();
			} catch {
				ctx.ui.notify("Failed to save plan.", "error");
			}
		},
	});

	pi.registerShortcut(Key.ctrlAlt("p"), {
		description: "Toggle plan mode",
		handler: async (ctx) => togglePlanMode(ctx),
	});

	// ---- submit_plan tool ----
	pi.registerTool({
		name: "submit_plan",
		label: "Submit Plan",
		description:
			"Submit a structured implementation plan for user approval. ONLY available in plan mode. " +
			"Call this once analysis is finished and you are ready to propose the work. After calling it, stop and wait for the user.",
		promptSnippet: "Submit a structured plan (title + numbered steps) for approval while in plan mode",
		promptGuidelines: [
			"Use submit_plan to deliver your plan when in plan mode instead of writing a free-form 'Plan:' section.",
			"After calling submit_plan, do not take further action; wait for the user to approve, refine, or execute.",
		],
		parameters: submitPlanSchema,
		async execute(_toolCallId, params: SubmitPlanInput, _signal, _onUpdate, ctx) {
			if (!planModeEnabled) {
				return {
					content: [
						{ type: "text", text: "submit_plan is only available in plan mode. Enable it with /plan first." },
					],
					isError: true,
					details: {},
				};
			}

			plan = {
				title: params.title.trim() || "Untitled plan",
				createdAt: new Date().toISOString(),
				steps: params.steps
					.map((s) => ({ text: s.description.trim(), completed: false }))
					.filter((s) => s.text.length > 0),
				risks: params.risks?.map((r) => r.trim()).filter(Boolean),
				files: params.files?.map((f) => f.trim()).filter(Boolean),
			};
			currentStepIndex = 0;
			awaitingStep = false;
			phase = "awaiting-approval";

			persistState();
			updateStatus(ctx);

			return {
				content: [
					{
						type: "text",
						text: `Plan "${plan.title}" with ${plan.steps.length} step(s) submitted. Stop here and wait for the user to approve, refine, or execute. Use /plan-save to save the plan to a file.`,
					},
				],
				details: { title: plan.title, steps: plan.steps.length },
				terminate: true,
			};
		},
	});

	// ---- block destructive bash in plan mode ----
	pi.on("tool_call", async (event) => {
		if (!planModeEnabled || event.toolName !== "bash") return;
		const command = event.input.command as string;
		if (!isSafeCommand(command)) {
			return {
				block: true,
				reason: `Plan mode: command blocked (not allowlisted). Use /plan to disable plan mode first.\nCommand: ${command}`,
			};
		}
	});

	// ---- strip stale plan-mode context when not in plan mode ----
	pi.on("context", async (event) => {
		if (planModeEnabled) return;
		return {
			messages: event.messages.filter((m) => {
				const msg = m as AgentMessage & { customType?: string };
				if (msg.customType === "plan-mode-context") return false;
				if (msg.role !== "user") return true;
				const content = msg.content;
				if (typeof content === "string") return !content.includes("[PLAN MODE ACTIVE]");
				if (Array.isArray(content)) {
					return !content.some(
						(c) => c.type === "text" && (c as TextContent).text?.includes("[PLAN MODE ACTIVE]"),
					);
				}
				return true;
			}),
		};
	});

	// ---- inject plan-mode instructions ----
	pi.on("before_agent_start", async () => {
		if (!planModeEnabled) return;
		return {
			message: {
				customType: "plan-mode-context",
				content: `[PLAN MODE ACTIVE]
You are in plan mode - a read-only exploration mode for safe code analysis.

Restrictions:
- Built-in edit and write tools are disabled
- Bash is restricted to an allowlist of read-only commands
- Do NOT attempt to make changes - only investigate and plan

Explore the code, ask clarifying questions if needed, then call the "submit_plan" tool with:
- title: a short, recognizable name for the plan
- steps: an ordered list of concrete, actionable steps
- risks: optional caveats or open questions
- files: optional files you expect to create or modify

Do NOT write a free-form "Plan:" section - use the submit_plan tool instead.
After submitting, stop and wait for the user.`,
				display: false,
			},
		};
	});

	// ---- drive approval + gated execution when the agent settles ----
	pi.on("agent_settled", async (_event, ctx) => {
		if (!ctx.hasUI) return;

		if (phase === "awaiting-approval") {
			await handleApproval(ctx);
			return;
		}

		if (phase === "executing" && awaitingStep) {
			awaitingStep = false;
			const step = plan?.steps[currentStepIndex];
			if (step) step.completed = true;
			currentStepIndex++;
			updateStatus(ctx);
			persistState();
			await persistPlanFile();
			await gateNext(ctx);
		}
	});

	// ---- restore state on session start / resume ----
	pi.on("session_start", async (_event, ctx) => {
		if (pi.getFlag("plan") === true) {
			planModeEnabled = true;
			phase = "planning";
		}

		const entries = ctx.sessionManager.getEntries();
		const planModeEntry = entries
			.filter((e: { type: string; customType?: string }) => e.type === "custom" && e.customType === "plan-mode")
			.pop() as { data?: PlanModeState } | undefined;

		if (planModeEntry?.data) {
			const d = planModeEntry.data;
			planModeEnabled = d.enabled ?? planModeEnabled;
			phase = d.phase ?? phase;
			plan = d.plan ?? plan;
			currentStepIndex = d.currentStepIndex ?? currentStepIndex;
			planFilePath = d.planFilePath ?? planFilePath;
			toolsBeforePlanMode = d.toolsBeforePlanMode ?? toolsBeforePlanMode;
			awaitingStep = false; // never resume mid-step automatically
		}

		if (planModeEnabled) {
			enablePlanModeTools();
		}
		updateStatus(ctx);

		if (phase === "executing" && plan && ctx.hasUI) {
			ctx.ui.notify(`Plan "${plan.title}" execution paused. Use /plan-continue to resume.`, "info");
		}
	});
}