import { spawn } from "node:child_process";
import * as fs from "node:fs";
import * as path from "node:path";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import {
	createBashTool,
	createEditTool,
	createReadTool,
	createWriteTool,
	type BashOperations,
	type EditOperations,
	type ReadOperations,
	type WriteOperations,
} from "@mariozechner/pi-coding-agent";

type HostConfig = {
	remote: string;
	remoteCwd?: string;
	description?: string;
	roles?: string[];
	allowWrite?: boolean;
	allowBash?: boolean;
	confirmDestructive?: boolean;
};

type SshConfig = {
	allowRawHosts?: boolean;
	defaultConfirmDestructive?: boolean;
	hosts?: Record<string, HostConfig>;
};

type ResolvedSsh = {
	alias?: string;
	remote: string;
	remoteCwd: string;
	allowWrite: boolean;
	allowBash: boolean;
	confirmDestructive: boolean;
	description?: string;
	roles: string[];
};

const CONFIG_PATH = path.join(process.cwd(), ".pi", "ssh", "hosts.json");
const SSH_OPTIONS = ["-o", "BatchMode=yes", "-o", "ConnectTimeout=10"];
const DANGEROUS_PATTERNS = [
	/\brm\s+(-rf?|--recursive)\b/i,
	/\bmkfs\b/i,
	/\bdd\b.+\bof=\/dev\//i,
	/\bshutdown\b/i,
	/\breboot\b/i,
	/\bpoweroff\b/i,
	/\bsystemctl\s+(reboot|poweroff)\b/i,
	/\bchmod\b.+\b777\b/i,
	/\bchown\b/i,
	/\bmv\b.+\s\/etc\//i,
	/\biptables\b/i,
	/\bnft\b/i,
	/\buserdel\b/i,
	/\bgroupdel\b/i,
];

function readConfig(): SshConfig {
	if (!fs.existsSync(CONFIG_PATH)) return {};
	const raw = fs.readFileSync(CONFIG_PATH, "utf8");
	return JSON.parse(raw) as SshConfig;
}

function parseSshArg(arg: string): { target: string; remoteCwd?: string } {
	const slashColon = arg.indexOf(":/");
	if (slashColon >= 0) {
		return {
			target: arg.slice(0, slashColon),
			remoteCwd: arg.slice(slashColon + 1),
		};
	}
	return { target: arg };
}

function shellQuote(value: string): string {
	return JSON.stringify(value);
}

function sshExec(remote: string, command: string): Promise<Buffer> {
	return new Promise((resolve, reject) => {
		const child = spawn("ssh", [...SSH_OPTIONS, remote, command], {
			stdio: ["ignore", "pipe", "pipe"],
		});
		const chunks: Buffer[] = [];
		const errChunks: Buffer[] = [];
		child.stdout.on("data", (data) => chunks.push(data));
		child.stderr.on("data", (data) => errChunks.push(data));
		child.on("error", reject);
		child.on("close", (code) => {
			if (code !== 0) {
				reject(new Error(`SSH failed (${code}): ${Buffer.concat(errChunks).toString().trim()}`));
				return;
			}
			resolve(Buffer.concat(chunks));
		});
	});
}

function resolveSshTarget(arg: string, config: SshConfig): ResolvedSsh {
	const { target, remoteCwd } = parseSshArg(arg);
	const host = config.hosts?.[target];
	if (host) {
		return {
			alias: target,
			remote: host.remote,
			remoteCwd: remoteCwd ?? host.remoteCwd ?? ".",
			allowWrite: host.allowWrite ?? true,
			allowBash: host.allowBash ?? true,
			confirmDestructive: host.confirmDestructive ?? config.defaultConfirmDestructive ?? true,
			description: host.description,
			roles: host.roles ?? [],
		};
	}

	if (!config.allowRawHosts) {
		throw new Error(
			`SSH target ${target} is not in .pi/ssh/hosts.json and allowRawHosts is disabled`,
		);
	}

	return {
		remote: target,
		remoteCwd: remoteCwd ?? ".",
		allowWrite: true,
		allowBash: true,
		confirmDestructive: config.defaultConfirmDestructive ?? true,
		roles: [],
	};
}

function normalizeRemotePath(remoteCwd: string, localCwd: string, inputPath: string): string {
	const absoluteLocal = path.isAbsolute(inputPath) ? path.resolve(inputPath) : path.resolve(localCwd, inputPath);
	const relative = path.relative(localCwd, absoluteLocal);
	if (relative.startsWith("..") || path.isAbsolute(relative)) {
		throw new Error(`Refusing to map path outside project root over SSH: ${inputPath}`);
	}
	const segments = relative.split(path.sep).filter(Boolean);
	const remoteBase = remoteCwd === "." ? "." : path.posix.normalize(remoteCwd);
	return segments.length === 0 ? remoteBase : path.posix.join(remoteBase, ...segments);
}

function isDangerousCommand(command: string): boolean {
	return DANGEROUS_PATTERNS.some((pattern) => pattern.test(command));
}

function createRemoteReadOps(remote: string, remoteCwd: string, localCwd: string): ReadOperations {
	const toRemote = (p: string) => normalizeRemotePath(remoteCwd, localCwd, p);
	return {
		readFile: (p) => sshExec(remote, `cat ${shellQuote(toRemote(p))}`),
		access: (p) => sshExec(remote, `test -r ${shellQuote(toRemote(p))}`).then(() => {}),
		detectImageMimeType: async (p) => {
			try {
				const result = await sshExec(remote, `file --mime-type -b ${shellQuote(toRemote(p))}`);
				const mime = result.toString().trim();
				return ["image/jpeg", "image/png", "image/gif", "image/webp"].includes(mime) ? mime : null;
			} catch {
				return null;
			}
		},
	};
}

function createRemoteWriteOps(remote: string, remoteCwd: string, localCwd: string): WriteOperations {
	const toRemote = (p: string) => normalizeRemotePath(remoteCwd, localCwd, p);
	return {
		writeFile: async (p, content) => {
			const encoded = Buffer.from(content).toString("base64");
			await sshExec(
				remote,
				`mkdir -p $(dirname ${shellQuote(toRemote(p))}) && echo ${shellQuote(encoded)} | base64 -d > ${shellQuote(toRemote(p))}`,
			);
		},
		mkdir: (dir) => sshExec(remote, `mkdir -p ${shellQuote(toRemote(dir))}`).then(() => {}),
	};
}

function createRemoteEditOps(remote: string, remoteCwd: string, localCwd: string): EditOperations {
	const readOps = createRemoteReadOps(remote, remoteCwd, localCwd);
	const writeOps = createRemoteWriteOps(remote, remoteCwd, localCwd);
	return {
		readFile: readOps.readFile,
		access: readOps.access,
		writeFile: writeOps.writeFile,
	};
}

function createRemoteBashOps(remote: string, remoteCwd: string, localCwd: string): BashOperations {
	return {
		exec: (command, cwd, { onData, signal, timeout }) =>
			new Promise((resolve, reject) => {
				const remoteDir = normalizeRemotePath(remoteCwd, localCwd, cwd);
				const remoteCommand = `cd ${shellQuote(remoteDir)} && ${command}`;
				const child = spawn("ssh", [...SSH_OPTIONS, remote, remoteCommand], {
					stdio: ["ignore", "pipe", "pipe"],
				});
				let timedOut = false;
				const timer = timeout
					? setTimeout(() => {
							timedOut = true;
							child.kill();
						}, timeout * 1000)
					: undefined;

				child.stdout.on("data", onData);
				child.stderr.on("data", onData);
				child.on("error", (error) => {
					if (timer) clearTimeout(timer);
					reject(error);
				});

				const onAbort = () => child.kill();
				signal?.addEventListener("abort", onAbort, { once: true });
				child.on("close", (code) => {
					if (timer) clearTimeout(timer);
					signal?.removeEventListener("abort", onAbort);
					if (signal?.aborted) reject(new Error("aborted"));
					else if (timedOut) reject(new Error(`timeout:${timeout}`));
					else resolve({ exitCode: code });
				});
			}),
	};
}

export default function (pi: ExtensionAPI) {
	pi.registerFlag("ssh", {
		description: "SSH host alias or user@host, optionally with :/remote/path",
		type: "string",
	});

	const localCwd = process.cwd();
	const localRead = createReadTool(localCwd);
	const localWrite = createWriteTool(localCwd);
	const localEdit = createEditTool(localCwd);
	const localBash = createBashTool(localCwd);

	let resolvedSsh: ResolvedSsh | null = null;
	const getSsh = () => resolvedSsh;

	pi.registerTool({
		...localRead,
		async execute(id, params, signal, onUpdate) {
			const ssh = getSsh();
			if (!ssh) return localRead.execute(id, params, signal, onUpdate);
			const tool = createReadTool(localCwd, {
				operations: createRemoteReadOps(ssh.remote, ssh.remoteCwd, localCwd),
			});
			return tool.execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...localWrite,
		async execute(id, params, signal, onUpdate) {
			const ssh = getSsh();
			if (!ssh) return localWrite.execute(id, params, signal, onUpdate);
			if (!ssh.allowWrite) throw new Error(`SSH target ${ssh.alias ?? ssh.remote} is read-only`);
			const tool = createWriteTool(localCwd, {
				operations: createRemoteWriteOps(ssh.remote, ssh.remoteCwd, localCwd),
			});
			return tool.execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...localEdit,
		async execute(id, params, signal, onUpdate) {
			const ssh = getSsh();
			if (!ssh) return localEdit.execute(id, params, signal, onUpdate);
			if (!ssh.allowWrite) throw new Error(`SSH target ${ssh.alias ?? ssh.remote} is read-only`);
			const tool = createEditTool(localCwd, {
				operations: createRemoteEditOps(ssh.remote, ssh.remoteCwd, localCwd),
			});
			return tool.execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...localBash,
		async execute(id, params, signal, onUpdate) {
			const ssh = getSsh();
			if (!ssh) return localBash.execute(id, params, signal, onUpdate);
			if (!ssh.allowBash) throw new Error(`SSH target ${ssh.alias ?? ssh.remote} does not allow bash execution`);
			const tool = createBashTool(localCwd, {
				operations: createRemoteBashOps(ssh.remote, ssh.remoteCwd, localCwd),
			});
			return tool.execute(id, params, signal, onUpdate);
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		const arg = pi.getFlag("ssh") as string | undefined;
		if (!arg) return;

		const config = readConfig();
		resolvedSsh = resolveSshTarget(arg, config);
		if (resolvedSsh.remoteCwd === ".") {
			resolvedSsh.remoteCwd = (await sshExec(resolvedSsh.remote, "pwd")).toString().trim();
		}

		if (ctx.hasUI) {
			const label = resolvedSsh.alias ?? resolvedSsh.remote;
			ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH: ${label}:${resolvedSsh.remoteCwd}`));
			ctx.ui.notify(`SSH mode: ${label}:${resolvedSsh.remoteCwd}`, "info");
		}
	});

	pi.on("tool_call", async (event, ctx) => {
		const ssh = getSsh();
		if (!ssh || event.toolName !== "bash") return;
		const command = event.input.command as string;
		if (!ssh.allowBash) {
			return { block: true, reason: `SSH bash disabled for ${ssh.alias ?? ssh.remote}` };
		}
		if (!ssh.confirmDestructive || !isDangerousCommand(command)) return;
		if (!ctx.hasUI) {
			return { block: true, reason: "Dangerous SSH command blocked (no UI for confirmation)" };
		}
		const choice = await ctx.ui.select(
			`⚠️ Dangerous SSH command on ${ssh.alias ?? ssh.remote}:\n\n${command}\n\nAllow?`,
			["Yes", "No"],
		);
		if (choice !== "Yes") {
			return { block: true, reason: "Blocked by user" };
		}
	});

	pi.on("user_bash", async (event, ctx) => {
		const ssh = getSsh();
		if (!ssh) return;
		if (!ssh.allowBash) {
			return { result: { output: `SSH bash disabled for ${ssh.alias ?? ssh.remote}`, exitCode: 1, cancelled: false, truncated: false } };
		}
		if (ssh.confirmDestructive && isDangerousCommand(event.command)) {
			if (!ctx.hasUI) {
				return { result: { output: "Dangerous SSH command blocked (no UI for confirmation)", exitCode: 1, cancelled: false, truncated: false } };
			}
			const ok = await ctx.ui.confirm("Dangerous SSH command", `${event.command}\n\nRun on ${ssh.alias ?? ssh.remote}?`);
			if (!ok) {
				return { result: { output: "Blocked by user", exitCode: 1, cancelled: false, truncated: false } };
			}
		}
		return { operations: createRemoteBashOps(ssh.remote, ssh.remoteCwd, localCwd) };
	});

	pi.on("before_agent_start", async (event) => {
		const ssh = getSsh();
		if (!ssh) return;
		const label = ssh.alias ?? ssh.remote;
		return {
			systemPrompt:
				event.systemPrompt.replace(
					`Current working directory: ${localCwd}`,
					`Current working directory: ${ssh.remoteCwd} (via SSH: ${label})`,
				) +
				`\n\nSSH target details:\n- target: ${label}\n- remote: ${ssh.remote}\n- remote cwd: ${ssh.remoteCwd}\n- roles: ${ssh.roles.join(", ") || "none"}\n- write access: ${ssh.allowWrite ? "enabled" : "disabled"}\n- bash access: ${ssh.allowBash ? "enabled" : "disabled"}`,
		};
	});
}