import type { ExtensionAPI, ExtensionCommandContext } from "@earendil-works/pi-coding-agent";
import { existsSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
import { execSync } from "node:child_process";
import { isAbsolute, join, relative, resolve } from "node:path";

const MODES = ["read-only", "edit"] as const;
type Mode = (typeof MODES)[number];

const WORKSPACE_TEMPLATE = `# Plan

<!-- What problem does this branch solve, and what is the approach? -->

# Discussion

<!-- Notes, tradeoffs, and decisions made. -->

# Progress

Not started.

# Next

<!-- Next steps and open tasks. -->
`;

const AGENTS_TEMPLATE = "# Agent Instructions\n\n<!-- Shared behavior and constraints for this repository. -->\n";

// =============================================================================
// Git / workspace helpers
// =============================================================================

function branch(cwd: string): string {
	try {
		return execSync("git symbolic-ref --quiet --short HEAD", { cwd, encoding: "utf8", stdio: ["ignore", "pipe", "ignore"] }).trim() || "detached";
	} catch { return "main"; }
}

function workspaceDir(cwd: string): string {
	return join(cwd, ".pi", branch(cwd).replace(/[\\/]/g, "-"));
}

function writeMissing(path: string, content: string): void {
	if (!existsSync(path)) writeFileSync(path, content, "utf8");
}

function read(path: string): string {
	return readFileSync(path, "utf8").trimEnd();
}

const MODE_ALIASES: Record<string, Mode> = { e: "edit", r: "read-only", discuss: "read-only", implement: "edit" };

function normalizeMode(value: string): Mode {
	const m = value.trim();
	if (m in MODE_ALIASES) return MODE_ALIASES[m]!;
	return (MODES as readonly string[]).includes(m) ? (m as Mode) : "read-only";
}

/** Parse H1 section names from workspace.md */
function parseHeaders(wsFile: string): string[] {
	try {
		return [...readFileSync(wsFile, "utf8").matchAll(/^# (.+)$/gm)].map(m => m[1]!.trim());
	} catch { return []; }
}

/** Migrate old plan/discuss/implementation.md → workspace.md */
function migrateWorkspace(dir: string): void {
	const wsFile = join(dir, "workspace.md");
	if (existsSync(wsFile)) return;
	const OLD = [
		{ file: "plan.md", header: "# Plan" },
		{ file: "discuss.md", header: "# Discussion" },
		{ file: "implementation.md", header: "# Progress" },
	];
	const parts: string[] = [];
	for (const { file, header } of OLD) {
		const p = join(dir, file);
		if (!existsSync(p)) continue;
		const body = readFileSync(p, "utf8").trim();
		if (body && body !== header) parts.push(body);
	}
	writeFileSync(wsFile, parts.length ? parts.join("\n\n") + "\n" : WORKSPACE_TEMPLATE, "utf8");
}

function ensure(cwd: string): { dir: string; wsFile: string; mode: Mode; reset: boolean } {
	const dir = workspaceDir(cwd);
	mkdirSync(dir, { recursive: true });
	writeMissing(join(cwd, ".pi", "agents.md"), AGENTS_TEMPLATE);
	migrateWorkspace(dir);
	const wsFile = join(dir, "workspace.md");

	const modePath = join(dir, ".mode");
	let reset = false;
	if (!existsSync(modePath)) writeFileSync(modePath, "read-only\n", "utf8");
	const raw = readFileSync(modePath, "utf8");
	const mode = normalizeMode(raw);
	if (mode !== raw.trim()) { writeFileSync(modePath, "read-only\n", "utf8"); reset = true; }
	return { dir, wsFile, mode, reset };
}

function isWorkspacePath(cwd: string, filePath: string | undefined): boolean {
	if (!filePath) return false;
	const pathFromPi = relative(resolve(cwd, ".pi"), resolve(cwd, filePath));
	return !pathFromPi.startsWith("..") && !isAbsolute(pathFromPi);
}

function isModeFile(cwd: string, filePath: string | undefined): boolean {
	if (!filePath) return false;
	return resolve(cwd, filePath) === resolve(join(workspaceDir(cwd), ".mode"));
}

function writesModeViaBash(command: string): boolean {
	const c = command.replace(/\d?>>?\s*\/dev\/null/g, "");
	return /\.mode\b/.test(c) && /(>+|tee\b|sed\s+-i|perl\s+-pi|mv\b|cp\b|rm\b)/.test(c);
}

// =============================================================================
// Extension
// =============================================================================

export default function (pi: ExtensionAPI) {
	// On session start / reload: always drop back to read-only.
	pi.on("session_start", (_event, ctx) => {
		const s = ensure(ctx.cwd);
		if (s.mode !== "read-only") writeFileSync(join(s.dir, ".mode"), "read-only\n", "utf8");
		if (ctx.hasUI && (s.reset || s.mode !== "read-only")) ctx.ui.notify("mode: read-only", "info");
	});

	// Inject branch workspace context into every agent turn.
	pi.on("before_agent_start", (event, ctx) => {
		const s = ensure(ctx.cwd);
		const agentsMd = existsSync(join(ctx.cwd, "AGENTS.md")) ? join(ctx.cwd, "AGENTS.md") : join(ctx.cwd, ".pi", "agents.md");
		const guard = s.mode === "read-only"
			? `\n\nMODE GUARD: mode is read-only — maximize reading, minimize output. ALWAYS read actual files first (use read/grep/bash). Trace real code paths. Verify every assumption with evidence from the codebase. Never propose a change before you have read the relevant files. Never guess. Never suggest an edit unprompted. If ambiguous after reading, ask one focused question. Do not write or edit project files outside .pi/**. You MAY write under .pi/** (${join(ctx.cwd, ".pi")}). Reads, chmod, and operations outside this project are allowed.`
			: "";
		return {
			systemPrompt: `${event.systemPrompt}\n\n## Per-Branch Workspace\nBranch: ${branch(ctx.cwd)}\nWorkspace: ${s.dir}\nCurrent mode: ${s.mode}${guard}\n\n### ${agentsMd.endsWith("AGENTS.md") ? "AGENTS.md" : ".pi/agents.md"}\n${read(agentsMd)}\n\n### workspace.md\n${read(s.wsFile)}`,
		};
	});

	// Gate tool calls based on current mode.
	pi.on("tool_call", (event, ctx) => {
		const s = ensure(ctx.cwd);
		const input = event.input as { path?: string; command?: string };

		// Always block direct mode-file manipulation
		if ((event.toolName === "write" || event.toolName === "edit") && isModeFile(ctx.cwd, input.path))
			return { block: true, reason: "Use /mode read-only|edit to change the mode file." };
		if (event.toolName === "bash" && input.command && writesModeViaBash(input.command))
			return { block: true, reason: "Use /mode read-only|edit to change the mode file." };

		if (s.mode !== "read-only") return;

		// read-only: project file writes are limited to .pi/**
		if (event.toolName === "write" || event.toolName === "edit") {
			if (isWorkspacePath(ctx.cwd, input.path)) return;
			return { block: true, terminate: true, reason: "Blocked: in read-only mode, project file writes are limited to .pi/**." };
		}
	});

	// =============================================================================
	// Commands
	// =============================================================================

	const applyMode = async (next: Mode, s: ReturnType<typeof ensure>, ctx: ExtensionCommandContext) => {
		writeFileSync(join(s.dir, ".mode"), `${next}\n`, "utf8");
		ctx.ui.notify(`mode: ${next}`, "info");
		if (next !== "edit") return;
		await ctx.waitForIdle();
		pi.sendUserMessage(`Mode is now edit. Carry out what we agreed on above — follow the conversation and the Plan / Next sections of workspace.md (${s.wsFile}). If nothing concrete was agreed on yet, say so and stop.`);
	};

	pi.registerCommand("mode", {
		description: "Show or set workspace mode: read-only (r), edit (e)",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			const s = ensure(ctx.cwd);
			const raw = args.trim().split(/\s+/)[0] ?? "";

			if (!raw) {
				if (!ctx.hasUI) { ctx.ui.notify(`mode: ${s.mode}`, "info"); return; }
				const choice = await ctx.ui.select(`Current mode: ${s.mode}`, [...MODES]);
				if (!choice) return;
				await applyMode(choice as Mode, s, ctx);
				return;
			}

			const next: Mode | null = raw in MODE_ALIASES ? MODE_ALIASES[raw]! : (MODES as readonly string[]).includes(raw) ? (raw as Mode) : null;
			if (!next) {
				ctx.ui.notify("mode must be: read-only (r), edit (e)", "error");
				return;
			}
			await applyMode(next, s, ctx);
		},
	});

	pi.registerCommand("e", {
		description: "Shortcut: switch to edit mode",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			await applyMode("edit", ensure(ctx.cwd), ctx);
		},
	});

	pi.registerCommand("r", {
		description: "Shortcut: switch to read-only mode",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			await applyMode("read-only", ensure(ctx.cwd), ctx);
		},
	});

	pi.registerCommand("ws", {
		description: "Update workspace.md. /ws [section] to target a section, /ws to pick from menu.",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			const s = ensure(ctx.cwd);
			const parts = args.trim().split(/\s+/);
			let section = parts[0]?.toLowerCase() ?? "";
			const extra = parts.slice(1).join(" ").trim();

			// No args → show selector
			if (!section) {
				const headers = parseHeaders(s.wsFile);
				const SEP = "─────";
				const options = [...headers, SEP, "refine", "commit", "new section..."];
				const choice = await ctx.ui.select("Update workspace section:", options);
				if (!choice || choice === SEP) return;

				if (choice === "new section...") {
					const name = await ctx.ui.input("Section name:", "");
					if (!name?.trim()) return;
					section = name.trim().toLowerCase();
				} else {
					section = choice.toLowerCase();
				}
			}

			await ctx.waitForIdle();

			let prompt: string;
			if (section === "commit") {
				prompt = `Commit all uncommitted changes in this repo with bash: \`git add -A\` then \`git commit -m "wip: <summary>"\` where <summary> is your own one-line summary of what changed (run \`git status\` and \`git diff\` first to write it). Nothing else.${extra ? ` Context: ${extra}.` : ""}`;
			} else if (section === "refine") {
				prompt = `Refine the writing in workspace.md (${s.wsFile}) — fix grammar, improve clarity and conciseness. Do not change meaning or content. Preserve all # H1 headers exactly as-is.`;
			} else {
				const label = section.charAt(0).toUpperCase() + section.slice(1);
				const exists = parseHeaders(s.wsFile).some(h => h.toLowerCase() === section);
				const hint = extra ? ` Focus on this in particular: ${extra}.` : "";
				prompt = exists
					? `Based on our conversation, update the **# ${label}** section in workspace.md (${s.wsFile}). Keep other sections untouched. High-level and concise.${hint}`
					: `Append a new **# ${label}** section at the end of workspace.md (${s.wsFile}) based on our conversation. High-level and concise.${hint}`;
			}

			pi.sendUserMessage(prompt);
		},
	});
}