/**
 * SSH Remote Execution Example
 *
 * Demonstrates delegating tool operations to a remote machine via SSH.
 * When --ssh is provided, read/write/edit/bash run on the remote.
 *
 * Usage:
 *   atomic -e ./ssh.ts --ssh user@host
 *   atomic -e ./ssh.ts --ssh user@host:/remote/path
 *
 * Requirements:
 *   - SSH key-based auth (no password prompts)
 *   - bash on remote
 */

import { spawn } from "node:child_process";
import { dirname } from "node:path";
import type { ExtensionAPI } from "@bastani/atomic";
import {
	type BashOperations,
	createBashTool,
	createCodingTools,
	createEditTool,
	createReadTool,
	createWriteTool,
	type EditOperations,
	type ReadOperations,
	type WriteOperations,
} from "@bastani/atomic";

function sshExec(remote: string, command: string): Promise<Buffer> {
	return new Promise((resolve, reject) => {
		const child = spawn("ssh", [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()}`));
			} else {
				resolve(Buffer.concat(chunks));
			}
		});
	});
}

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

function createRemoteWriteOps(remote: string, remoteCwd: string, localCwd: string): WriteOperations {
	const toRemote = (p: string) => p.replace(localCwd, remoteCwd);
	return {
		writeFile: async (p, content, options) => {
			const b64 = Buffer.from(content).toString("base64");
			const target = JSON.stringify(toRemote(p));
			// `set -C` is the shell's noclobber, so `>` on an existing path fails instead of
			// truncating it. That is the remote equivalent of `O_EXCL`, and it is what keeps the
			// create path from silently taking a file that appeared since the read below.
			const redirect = options?.exclusive ? `set -C; base64 -d > ${target}` : `base64 -d > ${target}`;
			try {
				await sshExec(remote, `echo ${JSON.stringify(b64)} | { ${redirect}; }`);
			} catch (error) {
				// A remote shell reports noclobber as a bare non-zero exit, carrying no `code` to read.
				// `write` renders an exclusive-create failure as a `target_exists` conflict only when the
				// rejection carries `EEXIST`, so the collision has to be named here; otherwise this backend
				// alone surfaces an opaque SSH error where every other one reports a typed conflict.
				if (!options?.exclusive) throw error;
				// Re-probe rather than assume the failure was a collision. A permission error, a full disk,
				// or a dropped connection fail the same write, and each of those stays itself.
				const probe = await sshExec(
					remote,
					`{ test -e ${target} || test -L ${target}; } && echo yes || echo no`,
				).catch(() => undefined);
				if (probe?.toString().trim() !== "yes") throw error;
				throw Object.assign(new Error(`EEXIST: file already exists, open '${toRemote(p)}'`), { code: "EEXIST" });
			}
		},
		mkdir: (dir) => sshExec(remote, `mkdir -p ${JSON.stringify(toRemote(dir))}`).then(() => {}),
		// Reads the same filesystem the write lands on, which is the point: checking the local
		// disk would let `write` decide what to do about a remote file by looking at a different
		// machine. A missing path is `undefined`; anything else, including an unreadable one,
		// stays a failure.
		readFile: async (p) => {
			const target = JSON.stringify(toRemote(p));
			// `write` has already created the parent, so failure to enter it must remain
			// an error. A subshell preserves the cwd used to resolve relative targets.
			const parent = JSON.stringify(toRemote(dirname(p)));
			const probe = await sshExec(
				remote,
				`(cd ${parent}) && { { test -e ${target} || test -L ${target}; } && echo yes || echo no; }`,
			);
			if (probe.toString().trim() === "no") return undefined;
			return (await sshExec(remote, `cat ${target}`)).toString("utf8");
		},
	};
}

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

function createRemoteBashOps(remote: string, remoteCwd: string, localCwd: string): BashOperations {
	const toRemote = (p: string) => p.replace(localCwd, remoteCwd);
	return {
		exec: (command, cwd, { onData, signal, timeout }) =>
			new Promise((resolve, reject) => {
				const cmd = `cd ${JSON.stringify(toRemote(cwd))} && ${command}`;
				const child = spawn("ssh", [remote, cmd], { 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", (e) => {
					if (timer) clearTimeout(timer);
					reject(e);
				});
				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 remote: user@host or user@host:/path", type: "string" });

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

	// Resolved lazily on session_start (CLI flags not available during factory)
	let resolvedSsh: { remote: string; remoteCwd: string } | null = null;
	// The bundle shares one snapshot store across read/edit/write until the next session.
	let tools = createCodingTools(localCwd);

	const getSsh = () => resolvedSsh;

	pi.registerTool({
		...localRead,
		async execute(id, params, signal, onUpdate, _ctx) {
			return tools.find((tool) => tool.name === "read")!.execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...localWrite,
		async execute(id, params, signal, onUpdate, _ctx) {
			return tools.find((tool) => tool.name === "write")!.execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...localEdit,
		async execute(id, params, signal, onUpdate, _ctx) {
			return tools.find((tool) => tool.name === "edit")!.execute(id, params, signal, onUpdate);
		},
	});

	pi.registerTool({
		...localBash,
		async execute(id, params, signal, onUpdate, _ctx) {
			return tools.find((tool) => tool.name === "bash")!.execute(id, params, signal, onUpdate);
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		tools = createCodingTools(localCwd);
		resolvedSsh = null;
		// Resolve SSH config now that CLI flags are available
		const arg = pi.getFlag("ssh") as string | undefined;
		if (arg) {
			if (arg.includes(":")) {
				const [remote, path] = arg.split(":");
				resolvedSsh = { remote, remoteCwd: path };
			} else {
				// No path given, evaluate pwd on remote
				const remote = arg;
				const pwd = (await sshExec(remote, "pwd")).toString().trim();
				resolvedSsh = { remote, remoteCwd: pwd };
			}
			const { remote, remoteCwd } = resolvedSsh;
			tools = createCodingTools(localCwd, {
				read: { operations: createRemoteReadOps(remote, remoteCwd, localCwd) },
				write: { operations: createRemoteWriteOps(remote, remoteCwd, localCwd) },
				edit: { operations: createRemoteEditOps(remote, remoteCwd, localCwd) },
				bash: { operations: createRemoteBashOps(remote, remoteCwd, localCwd) },
			});
			ctx.ui.setStatus("ssh", ctx.ui.theme.fg("accent", `SSH: ${resolvedSsh.remote}:${resolvedSsh.remoteCwd}`));
			ctx.ui.notify(`SSH mode: ${resolvedSsh.remote}:${resolvedSsh.remoteCwd}`, "info");
		}
	});

	// Handle user ! commands via SSH
	pi.on("user_bash", (_event) => {
		const ssh = getSsh();
		if (!ssh) return; // No SSH, use local execution
		return { operations: createRemoteBashOps(ssh.remote, ssh.remoteCwd, localCwd) };
	});

	// Replace local cwd with remote cwd in system prompt
	pi.on("before_agent_start", async (event) => {
		const ssh = getSsh();
		if (ssh) {
			const modified = event.systemPrompt.replace(
				`Current working directory: ${localCwd}`,
				`Current working directory: ${ssh.remoteCwd} (via SSH: ${ssh.remote})`,
			);
			return { systemPrompt: modified };
		}
	});
}