/**
 * Git Checkout Guard Extension
 *
 * Prevents models from using `git checkout` or `git restore` to silently
 * discard uncommitted changes in files.
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { isToolCallEventType } from "@mariozechner/pi-coding-agent";
import { execSync } from "child_process";

/**
 * Parse file paths from a git checkout/restore command.
 * Returns null if the command doesn't look like a file-restore operation.
 */
function parseFileRestoreArgs(command: string): string[] | null {
	// Normalize whitespace
	const cmd = command.trim().replace(/\s+/g, " ");

	// Match: git checkout -- <files>
	// Match: git checkout <ref> -- <files>
	const checkoutDashDash = cmd.match(/\bgit\s+checkout\b.*?\s--\s+(.+)/);
	if (checkoutDashDash) {
		return checkoutDashDash[1].trim().split(/\s+/);
	}

	// Match: git restore [--staged] [--source=<ref>] <files>
	// (git restore always operates on files)
	const restore = cmd.match(/\bgit\s+restore\s+(.+)/);
	if (restore) {
		// Filter out flags like --staged, --source=..., --worktree, --patch
		const args = restore[1].trim().split(/\s+/);
		const files = args.filter((a) => !a.startsWith("-"));
		return files.length > 0 ? files : null;
	}

	return null;
}

/**
 * Check which of the given file paths have uncommitted changes (staged or unstaged).
 * Returns the subset that are dirty.
 */
function getDirtyFiles(files: string[], cwd: string): string[] {
	const dirty: string[] = [];
	for (const file of files) {
		try {
			// --porcelain output is empty for clean files
			const out = execSync(`git status --porcelain -- ${JSON.stringify(file)}`, {
				cwd,
				encoding: "utf8",
				stdio: ["ignore", "pipe", "ignore"],
			}).trim();
			if (out.length > 0) {
				dirty.push(file);
			}
		} catch {
			// Not a git repo or other error — skip
		}
	}
	return dirty;
}

export default function (pi: ExtensionAPI) {
	pi.on("tool_call", async (event, ctx) => {
		if (!isToolCallEventType("bash", event)) return undefined;

		const command: string = event.input.command ?? "";
		const files = parseFileRestoreArgs(command);
		if (!files || files.length === 0) return undefined;

		const cwd = process.cwd();
		const dirty = getDirtyFiles(files, cwd);
		if (dirty.length === 0) return undefined; // nothing to protect

		const fileList = dirty.map((f) => `  • ${f}`).join("\n");

		if (!ctx.hasUI) {
			return {
				block: true,
				reason: `git-checkout-guard: the following files have uncommitted changes and cannot be silently reverted:\n${fileList}\nShow the diff to the user and ask for explicit confirmation first.`,
			};
		}

		const choice = await ctx.ui.select(
			`⚠️  git-checkout-guard\n\nThe command:\n  ${command}\n\nwould discard uncommitted changes in:\n${fileList}\n\nProceed?`,
			["No, cancel (show diff instead)", "Yes, discard changes anyway"],
		);

		if (choice !== "Yes, discard changes anyway") {
			return {
				block: true,
				reason: `Blocked by git-checkout-guard. Run \`git diff ${dirty.join(" ")}\` and review before discarding.`,
			};
		}

		return undefined;
	});
}