// pi-guard: route the write-capable built-in tools through an OS sandbox
// (bubblewrap on Linux, sandbox-exec on macOS) so the OS — not a heuristic —
// enforces read-only / write-restricted execution on the *real* host filesystem.
// The host's own binaries and environment are shared with the sandboxed shell.
// Toggle in-session with /guard. See README.
//
// Modes:
//   off         - host tools (pi default)
//   readonly    - whole filesystem readable, every write fails at the syscall
//   restricted  - cwd (+ /guard allow dirs) writable; writes elsewhere fail
//
// Reads (read/ls/find/grep) run as pi's normal host tools — a read can't breach
// a *write* boundary. A real write boundary, not a hardened sandbox (reads are
// not confined, network is left on).

import { execSync } from "node:child_process";
import { createHash } from "node:crypto";
import { mkdtempSync, rmSync } from "node:fs";
import { homedir, tmpdir } from "node:os";
import { join, resolve as resolvePath } from "node:path";
import { Type } from "typebox";
import {
	createBashTool,
	createEditTool,
	createWriteTool,
	RpcClient,
	type ExtensionAPI,
	type ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type { Component, TUI } from "@earendil-works/pi-tui";
import { Container, Spacer, Text } from "@earendil-works/pi-tui";
import { type Mode, loadConfig } from "./config.ts";
import { bashOps, editOps, writeOps } from "./operations.ts";
import { type Backend, type SandboxMode, type SandboxSpec, Sandbox, detectBackend } from "./sandbox.ts";

// ---- Types and components for spawn_agents --------------------------------

const SPINNER_FRAMES = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"];
const SPAWN_WIDGET_KEY = "pi-guard-spawn";

type AgentState = "pending" | "starting" | "running" | "completed" | "error";

interface AgentStatus {
	index: number;
	prompt: string;
	state: AgentState;
	result?: string;
	error?: string;
	lastEvent?: string;
}

class AgentWidget implements Component {
	private agents: AgentStatus[];
	private container: Container;
	private cachedWidth?: number;
	private cachedLines?: string[];
	private spinnerFrame = 0;
	private spinnerTimer?: ReturnType<typeof setInterval>;

	constructor(
		private tui: TUI,
		private theme: any,
		prompts: string[],
	) {
		this.agents = prompts.map((prompt, index) => ({
			index,
			prompt: prompt.slice(0, 60) + (prompt.length > 60 ? "..." : ""),
			state: "pending" as AgentState,
		}));
		this.container = new Container();
		this.rebuild();
		this.spinnerTimer = setInterval(() => {
			this.spinnerFrame = (this.spinnerFrame + 1) % SPINNER_FRAMES.length;
			if (this.agents.some((a) => a.state === "running")) {
				this.rebuild();
				this.invalidate();
				this.tui.requestRender();
			}
		}, 130);
	}

	private rebuild(): void {
		this.container.clear();

		// Compact header
		this.container.addChild(
			new Text(
				this.theme.fg(
					"accent",
					this.theme.bold(`● Spawning ${this.agents.length} subagent${this.agents.length > 1 ? "s" : ""}`),
				),
				0,
				0,
			),
		);
		this.container.addChild(new Spacer(1));

		// Agent status lines
		for (const agent of this.agents) {
			const icon = this.getStateIcon(agent.state);
			const color = this.getStateColor(agent.state);
			const stateText = agent.state.toUpperCase().padEnd(10);

			const statusLine = `${icon} [${agent.index + 1}/${this.agents.length}] ${stateText} ${agent.prompt}`;
			this.container.addChild(new Text(this.theme.fg(color, statusLine), 1, 0));

			// Show last event if running
			if (agent.lastEvent) {
				const eventPreview = agent.lastEvent.slice(0, 70) + (agent.lastEvent.length > 70 ? "..." : "");
				this.container.addChild(new Text(this.theme.fg("dim", `    └─ ${eventPreview}`), 1, 0));
			}

			// Show error if failed
			if (agent.error) {
				const errorMsg = agent.error.slice(0, 70) + (agent.error.length > 70 ? "..." : "");
				this.container.addChild(new Text(this.theme.fg("error", `    └─ Error: ${errorMsg}`), 1, 0));
			}
		}

		this.container.addChild(new Spacer(1));

		// Summary line
		const completed = this.agents.filter((a) => a.state === "completed").length;
		const running = this.agents.filter((a) => a.state === "running" || a.state === "starting").length;
		const errored = this.agents.filter((a) => a.state === "error").length;

		this.container.addChild(
			new Text(
				this.theme.fg("muted", `  Completed: ${completed} | Running: ${running} | Failed: ${errored}`),
				0,
				0,
			),
		);

		// Help text
		const allDone = completed + errored === this.agents.length;
		if (!allDone) {
			this.container.addChild(new Text(this.theme.fg("dim", "  Working... (Ctrl+C to cancel)"), 0, 0));
		}
	}

	private getStateIcon(state: AgentState): string {
		switch (state) {
			case "pending":
				return "○";
			case "starting":
				return "◔";
			case "running":
				return SPINNER_FRAMES[this.spinnerFrame];
			case "completed":
				return "✓";
			case "error":
				return "✗";
		}
	}

	private getStateColor(state: AgentState): string {
		switch (state) {
			case "pending":
				return "dim";
			case "starting":
				return "muted";
			case "running":
				return "accent";
			case "completed":
				return "success";
			case "error":
				return "error";
		}
	}

	updateAgent(index: number, update: Partial<AgentStatus>): void {
		Object.assign(this.agents[index]!, update);
		// Rebuild the container with updated state
		this.rebuild();
		this.tui.requestRender();
	}

	render(width: number): string[] {
		if (this.cachedLines && this.cachedWidth === width) {
			return this.cachedLines;
		}
		this.cachedLines = this.container.render(width);
		this.cachedWidth = width;
		return this.cachedLines;
	}

	invalidate(): void {
		this.cachedWidth = undefined;
		this.cachedLines = undefined;
		this.container.invalidate();
	}

	dispose(): void {
		if (this.spinnerTimer) {
			clearInterval(this.spinnerTimer);
			this.spinnerTimer = undefined;
		}
	}
}

// ---- Extension entry point ------------------------------------------------

export default function (pi: ExtensionAPI) {
	const cwd = process.cwd();

	// Local (host) tool implementations used when mode === "off".
	const local = {
		write: createWriteTool(cwd),
		edit: createEditTool(cwd),
		bash: createBashTool(cwd),
	};

	const state = {
		cfg: loadConfig(cwd),
		mode: "off" as Mode,
		allow: [] as string[],
		backend: undefined as Backend | undefined,
		backendChecked: false,
		sandbox: undefined as Sandbox | undefined,
		sandboxSig: "",
		scratch: undefined as string | undefined,
		starting: undefined as Promise<Sandbox> | undefined,
	};
	state.mode = state.cfg.mode;
	state.allow = [...state.cfg.allow];

	pi.registerFlag("guard", {
		description: "Start pi-guard in a mode: off | readonly | restricted",
		type: "string",
		default: "",
	});

	pi.registerFlag("guard-allow", {
		description: "Additional writable directories for restricted mode (comma-separated)",
		type: "string",
		default: "",
	});

	function sessionInfo(ctx: ExtensionContext): { id: string; file: string } {
		const sm = ctx.sessionManager as unknown as {
			getSessionId?: () => string;
			getSessionFile?: () => string | undefined;
		};
		return { id: sm.getSessionId?.() ?? "", file: sm.getSessionFile?.() ?? "" };
	}

	async function ensureBackend(): Promise<Backend> {
		if (!state.backendChecked) {
			state.backend = await detectBackend(state.cfg.backend);
			state.backendChecked = true;
		}
		if (!state.backend) {
			const hint =
				process.platform === "linux"
					? "install bubblewrap (e.g. `apt install bubblewrap` / `dnf install bubblewrap`)"
					: process.platform === "darwin"
						? "sandbox-exec not found at /usr/bin/sandbox-exec"
						: "no supported sandbox backend on this platform (Linux and macOS only)";
			throw new Error(
				`pi-guard: no sandbox backend available — ${hint}. Install one, set "backend" in config, or run /guard off. Refusing to run on the host in ${state.mode} mode (fail-closed).`,
			);
		}
		return state.backend;
	}

	// One persistent scratch dir per session, reused across every tool call so
	// temp files created by one command are visible to the next (TMPDIR points
	// here). Created lazily; removed on session shutdown.
	// ponytail: OS temp cleaners reap any dir leaked by a hard crash; no GC sweep.
	function ensureScratch(): string {
		if (!state.scratch) state.scratch = mkdtempSync(join(tmpdir(), "pi-guard-"));
		return state.scratch;
	}

	function sandboxSignature(backend: Backend): string {
		return createHash("sha256")
			.update(JSON.stringify({ m: state.mode, c: cwd, a: [...state.allow].sort(), b: backend }))
			.digest("hex")
			.slice(0, 8);
	}

	async function ensureSandbox(ctx: ExtensionContext): Promise<Sandbox> {
		const backend = await ensureBackend();
		const sig = sandboxSignature(backend);
		if (state.sandbox && state.sandboxSig === sig) return state.sandbox;
		if (state.starting) return state.starting;

		const spec: SandboxSpec = {
			backend,
			mode: state.mode as SandboxMode,
			cwd,
			allow: state.allow,
			scratch: ensureScratch(),
		};
		state.starting = (async () => {
			const s = await Sandbox.create(spec);
			state.sandbox = s;
			state.sandboxSig = sig;
			refreshStatus(ctx);
			return s;
		})().finally(() => {
			state.starting = undefined;
		});
		return state.starting;
	}

	function cleanupScratch(): void {
		if (state.scratch) rmSync(state.scratch, { recursive: true, force: true });
		state.scratch = undefined;
		state.sandbox = undefined;
		state.sandboxSig = "";
		state.starting = undefined;
	}

	function refreshStatus(ctx: ExtensionContext): void {
		// ctx can go stale if this runs from a detached background continuation
		// (e.g. the session_start pre-warm below) after the session has already
		// ended — most visibly in print/one-shot mode, where the session can
		// shut down before sandbox creation (which spawns real subprocesses)
		// finishes. Swallow rather than crash the process; there's no status
		// bar left to update anyway.
		try {
			if (state.mode === "off") {
				ctx.ui.setStatus("pi-guard", "");
				return;
			}
			const be = state.backend ?? "no backend";
			const icon = state.mode === "readonly" ? "🔒" : "✎";
			ctx.ui.setStatus("pi-guard", `guard: ${state.mode} ${icon} [${be}]`);
		} catch {
			// stale ctx — nothing to update.
		}
	}

	// Re-apply state after a mode/allow change: rebuild or tear down the sandbox.
	async function applyChange(ctx: ExtensionContext): Promise<void> {
		try {
			if (state.mode === "off") {
				state.sandbox = undefined;
				state.sandboxSig = "";
				refreshStatus(ctx);
				return;
			}
			state.sandbox = undefined;
			state.sandboxSig = "";
			await ensureSandbox(ctx);
			refreshStatus(ctx);
		} catch (e) {
			refreshStatus(ctx);
			ctx.ui.notify(`pi-guard: ${(e as Error).message}`, "error");
		}
	}

	// ---- Tool overrides -------------------------------------------------------
	// Only write-capable tools are routed through the sandbox. read/ls/find/grep
	// stay as pi's built-in host tools (reads can't breach a write boundary).

	pi.registerTool({
		...local.write,
		async execute(id, params, signal, onUpdate, ctx) {
			if (state.mode === "off") return local.write.execute(id, params, signal, onUpdate);
			const s = await ensureSandbox(ctx);
			return createWriteTool(cwd, { operations: writeOps(s) }).execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...local.edit,
		async execute(id, params, signal, onUpdate, ctx) {
			if (state.mode === "off") return local.edit.execute(id, params, signal, onUpdate);
			const s = await ensureSandbox(ctx);
			return createEditTool(cwd, { operations: editOps(s) }).execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...local.bash,
		async execute(id, params, signal, onUpdate, ctx) {
			if (state.mode === "off") return local.bash.execute(id, params, signal, onUpdate);
			const s = await ensureSandbox(ctx);
			return createBashTool(cwd, { operations: bashOps(s) }).execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		name: "spawn_agents",
		description:
			"Spawn one or more subagents concurrently under the same guard mode. All run in parallel; returns when all are done, results in order. Shows real-time status of each agent.",
		parameters: Type.Object({
			prompts: Type.Array(Type.String(), {
				description: "Task prompts to run. Each prompt spawns one subagent; all run concurrently.",
				minItems: 1,
			}),
		}),
		async execute(_id, params: { prompts: string[] }, signal, _onUpdate, ctx) {
			const guardArgs: string[] = [];
			if (state.mode !== "off") {
				guardArgs.push("--guard", state.mode);
				if (state.allow.length > 0) {
					guardArgs.push("--guard-allow", state.allow.join(","));
				}
			}
			let piPath: string;
			try {
				piPath = execSync("which pi", { encoding: "utf8" }).trim();
				if (!piPath) throw new Error("pi binary not found in PATH");
			} catch (err) {
				throw new Error(`Failed to locate pi binary: ${err}`);
			}

			// If not in TUI mode, use simple implementation without UI
			if (ctx.mode !== "tui") {
				const results = await Promise.all(
					params.prompts.map(async (prompt) => {
						const client = new RpcClient({ cwd, args: guardArgs, cliPath: piPath });
						await client.start();
						try {
							const events = await client.promptAndWait(prompt, [], 300_000);
							return (
								events
									.findLast((e: any) => e.type === "message_end")
									?.message.content.findLast((c: any) => c.type === "text")?.text ?? "(no output)"
							);
						} finally {
							await client.stop();
						}
					}),
				);
				const outputText =
					results.length === 1 ? results[0] : results.map((r, i) => `[${i + 1}]\n${r}`).join("\n\n");
				return {
					content: [{ type: "text", text: outputText }],
				};
			}

			// TUI mode: non-blocking live status widget (setWidget), never steals focus.
			const results: (string | null)[] = new Array(params.prompts.length).fill(null);
			const clients: (RpcClient | null)[] = new Array(params.prompts.length).fill(null);
			let aborted = false;
			let widget: AgentWidget | undefined;

			const onAbort = () => {
				aborted = true;
				for (const client of clients) {
					if (!client) continue;
					client.abort().catch(() => {});
					client.stop().catch(() => {});
				}
			};

			try {
				ctx.ui.setWidget(
					SPAWN_WIDGET_KEY,
					(tui, theme) => {
						widget = new AgentWidget(tui, theme, params.prompts);
						return widget;
					},
					{ placement: "aboveEditor" },
				);
				ctx.ui.setStatus(
					SPAWN_WIDGET_KEY,
					`spawn_agents: running ${params.prompts.length} agent${params.prompts.length > 1 ? "s" : ""}…`,
				);
				signal?.addEventListener("abort", onAbort);

				const agentPromises = params.prompts.map(async (prompt, index) => {
					let unsubscribe: (() => void) | null = null;
					try {
						widget!.updateAgent(index, { state: "starting" });

						const client = new RpcClient({ cwd, args: guardArgs, cliPath: piPath });
						clients[index] = client;

						let finalResult = "(no output)";
						unsubscribe = client.onEvent((event: any) => {
							if (aborted) return;
							if (event.type === "tool_call") {
								widget!.updateAgent(index, { lastEvent: `Calling: ${event.toolCall?.name || "tool"}` });
							} else if (event.type === "tool_call_result") {
								const toolName = event.toolCall?.name || "tool";
								const success = !event.result?.isError;
								widget!.updateAgent(index, { lastEvent: success ? `✓ ${toolName}` : `✗ ${toolName} (error)` });
							} else if (event.type === "message_start") {
								widget!.updateAgent(index, { lastEvent: "Thinking..." });
							} else if (event.type === "message_delta") {
								widget!.updateAgent(index, { lastEvent: "Generating response..." });
							} else if (event.type === "message_end") {
								const text = event.message?.content?.findLast((c: any) => c.type === "text")?.text;
								if (text) finalResult = text;
							}
						});

						await client.start();
						widget!.updateAgent(index, { state: "running", lastEvent: "Starting agent..." });

						await client.prompt(prompt, []);
						await client.waitForIdle(300_000);

						unsubscribe();
						unsubscribe = null;

						if (aborted) {
							results[index] = "(aborted)";
							widget!.updateAgent(index, { state: "error", error: "Aborted by user" });
							return;
						}

						results[index] = finalResult;
						widget!.updateAgent(index, { state: "completed", result: finalResult });
						await client.stop();
					} catch (error) {
						if (unsubscribe) unsubscribe();
						if (aborted) {
							results[index] = "(aborted)";
							widget!.updateAgent(index, { state: "error", error: "Aborted" });
						} else {
							results[index] = `Error: ${error}`;
							widget!.updateAgent(index, { state: "error", error: String(error) });
						}
						await clients[index]?.stop().catch(() => {});
					}
				});

				await Promise.all(agentPromises);
			} finally {
				signal?.removeEventListener("abort", onAbort);
				widget?.dispose();
				ctx.ui.setWidget(SPAWN_WIDGET_KEY, undefined);
				ctx.ui.setStatus(SPAWN_WIDGET_KEY, undefined);
			}

			if (aborted) {
				return { content: [{ type: "text", text: "Aborted" }] };
			}
			const outputText =
				results.length === 1 ? results[0]! : results.map((r, i) => `[${i + 1}]\n${r}`).join("\n\n");
			return { content: [{ type: "text", text: outputText }] };
		},
	});

	pi.on("user_bash", async (_event, ctx) => {
		if (state.mode === "off") return;
		const s = await ensureSandbox(ctx);
		return { operations: bashOps(s) };
	});

	// ---- Lifecycle ------------------------------------------------------------

	pi.on("session_start", async (_event, ctx) => {
		state.cfg = loadConfig(ctx.cwd, ctx.isProjectTrusted());
		const flag = (pi.getFlag("guard") as string) || "";
		if (flag === "off" || flag === "readonly" || flag === "restricted") state.mode = flag;
		else state.mode = state.cfg.mode;
		state.allow = [...state.cfg.allow];
		const allowFlag = (pi.getFlag("guard-allow") as string) || "";
		if (allowFlag) {
			const additionalDirs = allowFlag.split(",").map(d => resolvePath(cwd, d.trim())).filter(Boolean);
			state.allow.push(...additionalDirs);
		}
		state.backend = undefined;
		state.backendChecked = false;

		refreshStatus(ctx);

		// Background: detect backend + pre-warm the sandbox. Detached from the
		// event handler (not awaited), so the session — and this ctx — may
		// already be gone by the time it settles (e.g. print/one-shot mode).
		// Never let that crash the process: swallow stale-ctx errors here too.
		void (async () => {
			try {
				await ensureBackend();
				if (state.mode !== "off") await ensureSandbox(ctx);
			} catch (e) {
				if (state.mode !== "off") {
					try {
						ctx.ui.notify(`pi-guard: ${(e as Error).message}`, "warning");
					} catch {
						// stale ctx — session already ended, nothing to notify.
					}
				}
			}
		})();
	});

	pi.on("session_shutdown", async (_event, _ctx) => {
		// No long-lived process to preserve; drop the scratch dir on any shutdown
		// (quit/new/resume/fork/reload). Temp files are session-scoped by design.
		cleanupScratch();
	});

	// Tell the model it operates under an OS write boundary on the real filesystem.
	pi.on("before_agent_start", async (event, _ctx) => {
		if (state.mode === "off") return;
		const be = state.backend ?? "sandbox";
		const note =
			state.mode === "readonly"
				? `Current working directory: ${cwd} (pi-guard ${be}, READ-ONLY: the filesystem is mounted read-only and write attempts fail at the syscall; your home directory is readable but not writable)`
				: `Current working directory: ${cwd} (pi-guard ${be}, restricted: writes under the working directory persist to the host; writes elsewhere — including your read-only home directory — fail)`;
		const hostLine = `Current working directory: ${cwd}`;
		const systemPrompt = event.systemPrompt.includes(hostLine)
			? event.systemPrompt.replace(hostLine, note)
			: `${event.systemPrompt}\n\n${note}`;
		return { systemPrompt };
	});

	// ---- Command --------------------------------------------------------------

	pi.registerCommand("guard", {
		description: "OS sandbox guard: /guard [readonly|restricted|off|allow <dir>|backend <name>|rebuild|status]",
		getArgumentCompletions: (prefix: string) => {
			const opts = ["readonly", "restricted", "off", "allow", "backend", "rebuild", "status"];
			const items = opts.filter((o) => o.startsWith(prefix)).map((o) => ({ value: o, label: o }));
			return items.length > 0 ? items : null;
		},
		handler: async (args, ctx) => {
			const [sub, ...rest] = args.trim().split(/\s+/).filter(Boolean);
			const arg = rest.join(" ").trim();

			switch (sub) {
				case undefined:
				case "":
				case "status": {
					await showStatus(ctx);
					break;
				}
				case "off":
					state.mode = "off";
					await applyChange(ctx);
					ctx.ui.notify("pi-guard: OFF (host tools).", "info");
					break;
				case "readonly":
				case "ro":
					state.mode = "readonly";
					ctx.ui.notify("pi-guard: switching to read-only…", "info");
					await applyChange(ctx);
					break;
				case "restricted":
				case "rw":
					state.mode = "restricted";
					ctx.ui.notify("pi-guard: switching to restricted…", "info");
					await applyChange(ctx);
					break;
				case "allow": {
					if (!arg) {
						ctx.ui.notify("usage: /guard allow <dir>", "warning");
						break;
					}
					const abs = resolvePath(cwd, arg);
					if (!state.allow.includes(abs)) state.allow.push(abs);
					if (state.mode === "off") state.mode = "restricted";
					ctx.ui.notify(`pi-guard: allow ${abs} (rebuilding sandbox)`, "info");
					await applyChange(ctx);
					break;
				}
				case "backend": {
					if (arg !== "bwrap" && arg !== "sandbox-exec" && arg !== "auto") {
						ctx.ui.notify("usage: /guard backend bwrap|sandbox-exec|auto", "warning");
						break;
					}
					state.cfg = { ...state.cfg, backend: arg };
					state.backendChecked = false;
					state.backend = undefined;
					ctx.ui.notify(`pi-guard: backend set to ${arg} (rebuilding)`, "info");
					await applyChange(ctx);
					break;
				}
				case "rebuild":
					ctx.ui.notify("pi-guard: rebuilding sandbox…", "info");
					await applyChange(ctx);
					ctx.ui.notify("pi-guard: sandbox rebuilt.", "info");
					break;
				default:
					ctx.ui.notify(`pi-guard: unknown subcommand '${sub}'.`, "warning");
			}
		},
	});

	async function showStatus(ctx: ExtensionContext): Promise<void> {
		let be = "unknown";
		try {
			be = await ensureBackend();
		} catch {
			be = "none available";
		}
		const { id, file } = sessionInfo(ctx);
		const lines = [
			`pi-guard status`,
			`  mode:      ${state.mode}`,
			`  backend:   ${be}`,
			`  workspace: ${cwd} (identity)`,
			`  home:      ${homedir()} (read-only)`,
			`  allow:     ${state.allow.length ? state.allow.join(", ") : "(cwd only)"}`,
			`  scratch:   ${state.scratch ?? "(none yet)"}`,
			`  session:   ${id || "(ephemeral)"}${file ? ` (${file})` : ""}`,
		];
		ctx.ui.notify(lines.join("\n"), "info");
	}
}