import { spawn } from "node:child_process";
import {
	DEFAULT_RLM_EXTRA_UV_ARGS,
	type KernelPythonSkill,
	resolveRuntimeIdentity,
	resolveRuntimeSourceDir,
} from "./bootstrap.js";
import { errorMessage } from "./shared.js";

const DEFAULT_REMOTE_CACHE_BASE = "$HOME/.cache/prime-agent/runtimes";
const STATE_SNAPSHOT_REQUIREMENT = "dill";

export interface RemoteHostConfig {
	/** Remote SSH destination (e.g. "chernobog" or "mayer@chernobog") */
	host: string;
	/** Path to an existing OpenSSH ControlMaster socket if any */
	controlPath?: string;
	/** Optional custom SSH port */
	port?: number;
	/** Optional identity file */
	identityFile?: string;
	/** Extra options to pass to SSH */
	sshOptions?: string[];
	/** Base directory for runtimes on the remote host. Defaults to "$HOME/.cache/prime-agent/runtimes" */
	cacheDir?: string;
	/** Optional explicit remote Python binary override */
	pythonCmd?: string;
}

export interface RemoteKernelRuntimeInfo {
	/** The content hash / identity of prime-agent-runtime */
	runtimeHash: string;
	/** Remote directory containing the cached runtime ($HOME/.cache/prime-agent/runtimes/<hash>) */
	remoteDir: string;
	/** Path to the python interpreter on the remote host ($HOME/.cache/prime-agent/runtimes/<hash>/.venv/bin/python) */
	remotePython: string;
	/** Whether this execution hit the cache (.ready was already present on the remote host) */
	cacheHit: boolean;
}

export interface EnsureRemoteKernelRuntimeOptions {
	hostConfig: RemoteHostConfig;
	pythonSkills?: readonly KernelPythonSkill[];
	onProgress?: (message: string) => void;
	signal?: AbortSignal;
}

export interface RemoteExecResult {
	exitCode: number;
	stdout: string;
	stderr: string;
}

/**
 * Builds the arguments array for the `ssh` command.
 */
export function buildSshArgs(hostConfig: RemoteHostConfig, remoteCommand?: string): string[] {
	const args: string[] = [];
	if (hostConfig.controlPath) {
		args.push("-S", hostConfig.controlPath);
	} else if (!hostConfig.sshOptions?.some((opt) => opt.includes("ControlMaster"))) {
		args.push("-o", "ControlMaster=auto", "-o", "ControlPersist=30m", "-o", "ControlPath=~/.ssh/cm-%C");
	}
	if (hostConfig.port !== undefined) {
		args.push("-p", String(hostConfig.port));
	}
	if (hostConfig.identityFile) {
		args.push("-i", hostConfig.identityFile);
	}
	if (hostConfig.sshOptions) {
		args.push(...hostConfig.sshOptions);
	}
	args.push(hostConfig.host);
	if (remoteCommand !== undefined) {
		args.push(remoteCommand);
	}
	return args;
}

/**
 * Computes the remote directory path keyed by content hash.
 */
export function getRemoteRuntimeDir(cacheBase: string | undefined, runtimeHash: string): string {
	const base = cacheBase ?? DEFAULT_REMOTE_CACHE_BASE;
	return `${base}/${runtimeHash}`;
}

/**
 * Executes a remote command via SSH and returns its exit code and output.
 */
export async function execRemote(
	hostConfig: RemoteHostConfig,
	remoteCommand: string,
	options: { signal?: AbortSignal } = {},
): Promise<RemoteExecResult> {
	const args = buildSshArgs(hostConfig, remoteCommand);
	return new Promise<RemoteExecResult>((resolve, reject) => {
		const child = spawn("ssh", args, {
			stdio: ["ignore", "pipe", "pipe"],
			signal: options.signal,
		});

		let stdout = "";
		let stderr = "";

		child.stdout?.on("data", (chunk: Buffer) => {
			stdout += chunk.toString("utf8");
		});
		child.stderr?.on("data", (chunk: Buffer) => {
			stderr += chunk.toString("utf8");
		});

		child.on("error", (err) => {
			reject(new Error(`Failed to execute SSH command on ${hostConfig.host}: ${errorMessage(err)}`));
		});

		child.on("close", (code) => {
			resolve({
				exitCode: code ?? 1,
				stdout,
				stderr,
			});
		});
	});
}

/**
 * Fast path: probes whether the content-addressed runtime already exists and is ready on the remote host.
 */
export async function probeRemoteKernelRuntime(
	hostConfig: RemoteHostConfig,
	remoteDir: string,
	signal?: AbortSignal,
): Promise<boolean> {
	const readyFile = `${remoteDir}/.ready`;
	const cmd = `test -f "${readyFile}"`;
	const result = await execRemote(hostConfig, cmd, { signal });
	return result.exitCode === 0;
}

/**
 * Streams a gzipped tarball of the local runtime source directory directly over SSH to the remote host.
 */
export async function transferRuntimeSource(
	hostConfig: RemoteHostConfig,
	sourceDir: string,
	remoteSrcDir: string,
	signal?: AbortSignal,
): Promise<void> {
	const sshExtractCmd = `mkdir -p "${remoteSrcDir}" && tar -xzf - -C "${remoteSrcDir}"`;
	const sshArgs = buildSshArgs(hostConfig, sshExtractCmd);

	const tarArgs = [
		"-czf",
		"-",
		"--exclude=.venv",
		"--exclude=__pycache__",
		"--exclude=.git",
		"--exclude=*.pyc",
		"--exclude=test",
		"-C",
		sourceDir,
		".",
	];

	return new Promise<void>((resolve, reject) => {
		const tarChild = spawn("tar", tarArgs, {
			stdio: ["ignore", "pipe", "pipe"],
			signal,
		});

		const sshChild = spawn("ssh", sshArgs, {
			stdio: ["pipe", "pipe", "pipe"],
			signal,
		});

		let sshStderr = "";
		let tarStderr = "";

		sshChild.stderr?.on("data", (chunk: Buffer) => {
			sshStderr += chunk.toString("utf8");
		});
		tarChild.stderr?.on("data", (chunk: Buffer) => {
			tarStderr += chunk.toString("utf8");
		});

		tarChild.on("error", (err) => {
			sshChild.kill();
			reject(new Error(`tar packaging failed: ${errorMessage(err)}`));
		});

		sshChild.on("error", (err) => {
			tarChild.kill();
			reject(new Error(`SSH transfer failed to ${hostConfig.host}: ${errorMessage(err)}`));
		});

		// Pipe tar stdout to ssh stdin
		tarChild.stdout?.pipe(sshChild.stdin);

		let tarClosed = false;
		let sshClosed = false;
		let tarExitCode: number | null = null;
		let sshExitCode: number | null = null;

		const checkDone = () => {
			if (!tarClosed || !sshClosed) return;
			if (tarExitCode !== 0) {
				reject(new Error(`tar command exited with code ${tarExitCode}: ${tarStderr.trim()}`));
			} else if (sshExitCode !== 0) {
				reject(new Error(`Remote extraction exited with code ${sshExitCode}: ${sshStderr.trim()}`));
			} else {
				resolve();
			}
		};

		tarChild.on("close", (code) => {
			tarClosed = true;
			tarExitCode = code;
			checkDone();
		});

		sshChild.on("close", (code) => {
			sshClosed = true;
			sshExitCode = code;
			checkDone();
		});
	});
}

/**
 * Builds the shell script to bootstrap a virtualenv and install dependencies on the remote host.
 */
export function buildRemoteBootstrapScript(
	remoteDir: string,
	extraUvArgs: readonly string[] = DEFAULT_RLM_EXTRA_UV_ARGS,
): string {
	const packages = ["-e src", STATE_SNAPSHOT_REQUIREMENT, ...extraUvArgs].join(" ");
	return `set -eu
TARGET="${remoteDir}"
cd "$TARGET"
if command -v uv >/dev/null 2>&1; then
    uv venv .venv
    uv pip install ${packages}
elif command -v nix >/dev/null 2>&1; then
    nix shell nixpkgs#uv -c uv venv .venv
    nix shell nixpkgs#uv -c uv pip install ${packages}
else
    python3 -m venv .venv
    .venv/bin/pip install ${packages}
fi
touch "$TARGET/.ready"
`;
}

/**
 * Ensures that a matching content-addressed prime-agent-runtime exists on the remote host.
 *
 * Fast path: probes for `$HOME/.cache/prime-agent/runtimes/<hash>/.ready`. If found, returns immediately.
 * Cold path: streams the runtime source tarball, creates a virtualenv, installs runtime dependencies,
 * and marks it ready.
 */
export async function ensureRemoteKernelRuntime(
	options: EnsureRemoteKernelRuntimeOptions,
): Promise<RemoteKernelRuntimeInfo> {
	const { hostConfig, onProgress, signal } = options;

	const remotePythonOverride = hostConfig.pythonCmd ?? process.env.PRIME_AGENT_REMOTE_PYTHON;
	if (remotePythonOverride) {
		return {
			runtimeHash: "custom",
			remoteDir: "",
			remotePython: remotePythonOverride,
			cacheHit: true,
		};
	}

	const runtimeHash = await resolveRuntimeIdentity();
	const remoteDir = getRemoteRuntimeDir(hostConfig.cacheDir, runtimeHash);
	const remotePython = `${remoteDir}/.venv/bin/python`;

	// Fast path: probe for .ready file
	const isReady = await probeRemoteKernelRuntime(hostConfig, remoteDir, signal);
	if (isReady) {
		if (options.pythonSkills && options.pythonSkills.length > 0) {
			await syncRemotePythonSkills(hostConfig, remoteDir, options.pythonSkills, signal, onProgress);
		}
		return {
			runtimeHash,
			remoteDir,
			remotePython,
			cacheHit: true,
		};
	}

	// Cold path: resolve local source directory
	const sourceDir = await resolveRuntimeSourceDir();
	if (!sourceDir) {
		throw new Error("Cannot find local prime-agent-runtime source directory to sync to remote host");
	}

	onProgress?.(`Transferring prime-agent-runtime (${runtimeHash.slice(0, 12)}) to ${hostConfig.host}...`);
	const remoteSrcDir = `${remoteDir}/src`;
	await transferRuntimeSource(hostConfig, sourceDir, remoteSrcDir, signal);

	onProgress?.(`Bootstrapping remote kernel environment on ${hostConfig.host}...`);
	const bootstrapScript = buildRemoteBootstrapScript(remoteDir);
	const bootstrapResult = await execRemote(hostConfig, bootstrapScript, { signal });
	if (bootstrapResult.exitCode !== 0) {
		throw new Error(
			`Failed to bootstrap remote kernel on ${hostConfig.host} (exit ${bootstrapResult.exitCode}): ` +
				`${bootstrapResult.stderr || bootstrapResult.stdout}`,
		);
	}

	if (options.pythonSkills && options.pythonSkills.length > 0) {
		await syncRemotePythonSkills(hostConfig, remoteDir, options.pythonSkills, signal, onProgress);
	}

	onProgress?.(`Remote kernel runtime ready on ${hostConfig.host}`);
	return {
		runtimeHash,
		remoteDir,
		remotePython,
		cacheHit: false,
	};
}

/**
 * Syncs user/workspace Python skills to the remote host and installs them in the remote kernel environment.
 */
export async function syncRemotePythonSkills(
	hostConfig: RemoteHostConfig,
	remoteDir: string,
	pythonSkills: readonly KernelPythonSkill[],
	signal?: AbortSignal,
	onProgress?: (message: string) => void,
): Promise<void> {
	if (!pythonSkills || pythonSkills.length === 0) return;

	for (const skill of pythonSkills) {
		const remoteSkillDir = `${remoteDir}/skills/${skill.importName}`;
		const probeCmd = `test -d "${remoteSkillDir}" && echo 1 || echo 0`;
		const probe = await execRemote(hostConfig, probeCmd, { signal });
		if (probe.exitCode === 0 && probe.stdout.trim() === "1") {
			continue;
		}

		onProgress?.(`Syncing Python skill ${skill.name || skill.importName} to ${hostConfig.host}...`);
		try {
			await transferRuntimeSource(hostConfig, skill.packagePath, remoteSkillDir, signal);
			const installCmd = `set -eu
TARGET="${remoteDir}"
SKILL="${remoteSkillDir}"
if command -v uv >/dev/null 2>&1; then
    uv pip install --python "$TARGET/.venv" -e "$SKILL"
elif command -v nix >/dev/null 2>&1; then
    nix shell nixpkgs#uv -c uv pip install --python "$TARGET/.venv" -e "$SKILL"
else
    "$TARGET/.venv/bin/pip" install -e "$SKILL"
fi
`;
			await execRemote(hostConfig, installCmd, { signal });
		} catch {
			// Failed skills are non-fatal; degraded skill handling is reported by kernel bootstrap
		}
	}
}