/**
 * Shell environment snapshot for preserving user aliases, functions, and options.
 *
 * Creates a snapshot file that captures the user's shell environment from their
 * .bashrc/.zshrc, which can be sourced before each command to provide a familiar
 * shell experience.
 */
import * as fs from "node:fs/promises";
import * as os from "node:os";
import * as path from "node:path";
import { postmortem } from "@gajae-code/utils";

const SNAPSHOT_TIMEOUT_MS = 2_000;
const SNAPSHOT_ROOT_PREFIX = "gjc-shell-snapshots-";
const PRIVATE_DIRECTORY_MODE = 0o700;
const PRIVATE_FILE_MODE = 0o600;

function sanitizeSnapshotEnv(env: Record<string, string | undefined>): Record<string, string | undefined> {
	const sanitized = { ...env };
	delete sanitized.BASH_ENV;
	delete sanitized.ENV;
	return sanitized;
}

/**
 * Get the user's shell config file path.
 */
function getShellConfigFile(shell: string, home: string): string {
	if (shell.includes("zsh")) return path.join(home, ".zshrc");
	if (shell.includes("bash")) return path.join(home, ".bashrc");
	return path.join(home, ".profile");
}

/**
 * Generate the snapshot creation script.
 * This script sources the user's rc file and extracts functions, aliases, and options.
 * Matches Anthropic Code's snapshot generation logic.
 */
function generateSnapshotScript(shell: string, snapshotPath: string, rcFile: string, hasRcFile: boolean): string {
	const isZsh = shell.includes("zsh");
	const commonToolsRegex =
		"^(ls|dir|vdir|cat|head|tail|less|more|grep|egrep|fgrep|rg|find|fd|locate|sed|awk|perl|cp|mv|rm|mkdir|rmdir|touch|chmod|chown|ln|pwd|readlink|stat|cut|sort|uniq|xargs|tee|tr|basename|dirname)$";

	// Escape the snapshot path for shell
	const escapedPath = snapshotPath.replace(/'/g, "'\\''");

	// Function extraction differs between bash and zsh
	const functionScript = isZsh
		? `
echo "# Functions" >> "$SNAPSHOT_FILE"
# Force autoload all functions first
typeset -f > /dev/null 2>&1
# Get user function names - filter system/private ones
typeset +f 2>/dev/null | grep -vE '^(_|__)' | grep -vE '${commonToolsRegex}' | while read func; do
   typeset -f "$func" >> "$SNAPSHOT_FILE" 2>/dev/null
done
`
		: `
echo "# Functions" >> "$SNAPSHOT_FILE"
# Force autoload all functions first
declare -f > /dev/null 2>&1
# Get user function names - filter system/private ones
declare -F 2>/dev/null | cut -d' ' -f3 | grep -vE '^(_|__)' | grep -vE '${commonToolsRegex}' | while read func; do
   declare -f "$func" >> "$SNAPSHOT_FILE" 2>/dev/null
done
`;

	// Shell options extraction
	const optionsScript = isZsh
		? `
echo "# Shell Options" >> "$SNAPSHOT_FILE"
setopt 2>/dev/null | sed 's/^/setopt /' | head -n 1000 >> "$SNAPSHOT_FILE"
`
		: `
echo "# Shell Options" >> "$SNAPSHOT_FILE"
shopt -p 2>/dev/null | head -n 1000 >> "$SNAPSHOT_FILE"
set -o 2>/dev/null | awk '$2 == "on" && $1 !~ /^(onecmd|monitor|restricted)$/ {print "set -o " $1}' | head -n 1000 >> "$SNAPSHOT_FILE"
echo "shopt -s expand_aliases" >> "$SNAPSHOT_FILE"
`;

	return `
umask 077
SNAPSHOT_FILE='${escapedPath}'

# Source user's rc file if it exists
${hasRcFile ? `source "${rcFile}" < /dev/null 2>/dev/null` : "# No user config file to source"}

# Create/clear the snapshot file
echo "# Shell snapshot - generated by gjc agent" >| "$SNAPSHOT_FILE"

# Unalias everything first to avoid conflicts when sourced
echo "unalias -a 2>/dev/null || true" >> "$SNAPSHOT_FILE"

${functionScript}

${optionsScript}

# Export aliases (limit to 1000)
echo "# Aliases" >> "$SNAPSHOT_FILE"
# Filter out winpty aliases on Windows to avoid "stdin is not a tty" errors
if [[ "$OSTYPE" == "msys" ]] || [[ "$OSTYPE" == "cygwin" ]]; then
   alias 2>/dev/null | grep -v "='winpty " | grep -vE '^alias (${commonToolsRegex})=' | sed 's/^alias //g' | sed 's/^/alias -- /' | head -n 1000 >> "$SNAPSHOT_FILE"
else
   alias 2>/dev/null | grep -vE '^alias (${commonToolsRegex})=' | sed 's/^alias //g' | sed 's/^/alias -- /' | head -n 1000 >> "$SNAPSHOT_FILE"
fi

# Export PATH
echo "export PATH='$PATH'" >> "$SNAPSHOT_FILE"

# Verify snapshot was created
if [ ! -f "$SNAPSHOT_FILE" ]; then
   echo "Error: Snapshot file was not created" >&2
   exit 1
fi
`.trim();
}

interface ShellSnapshotCacheOptions {
	tempRoot: string;
	home: string | (() => string);
	platform: NodeJS.Platform;
}

function ownedByCurrentUser(uid: number): boolean {
	return typeof process.getuid !== "function" || uid === process.getuid();
}

async function isTrustedSnapshot(snapshotPath: string): Promise<boolean> {
	try {
		const stat = await fs.lstat(snapshotPath);
		return stat.isFile() && !stat.isSymbolicLink() && ownedByCurrentUser(stat.uid) && (stat.mode & 0o077) === 0;
	} catch {
		return false;
	}
}

function createShellSnapshotCache(options: ShellSnapshotCacheOptions) {
	const cachedPaths = new Map<string, string>();
	const inFlight = new Map<string, Promise<string | null>>();
	let rootPath: string | null = null;
	let rootInitialization: Promise<string> | null = null;
	let cleanupAttempt: Promise<void> | null = null;
	let cleaningUp = false;
	let lifecycle = 0;

	async function initializeRoot(): Promise<string> {
		const root = await fs.mkdtemp(path.join(options.tempRoot, SNAPSHOT_ROOT_PREFIX));
		try {
			await fs.chmod(root, PRIVATE_DIRECTORY_MODE);
			const stat = await fs.lstat(root);
			if (
				!stat.isDirectory() ||
				stat.isSymbolicLink() ||
				!ownedByCurrentUser(stat.uid) ||
				(stat.mode & 0o777) !== 0o700
			) {
				throw new Error("Shell snapshot root is not owner-private");
			}
			return root;
		} catch (error) {
			await fs.rm(root, { recursive: true, force: true }).catch(() => {});
			throw error;
		}
	}

	async function getRoot(): Promise<string> {
		if (rootPath) return rootPath;
		if (rootInitialization) return await rootInitialization;
		const attempt = initializeRoot();
		rootInitialization = attempt;
		try {
			rootPath = await attempt;
			return rootPath;
		} catch (error) {
			if (rootInitialization === attempt) rootInitialization = null;
			throw error;
		}
	}

	async function createSnapshot(shell: string, env: Record<string, string | undefined>): Promise<string | null> {
		const root = await getRoot();
		const shellName = shell.includes("zsh") ? "zsh" : shell.includes("bash") ? "bash" : "sh";
		const snapshotPath = path.join(root, `snapshot-${shellName}-${crypto.randomUUID()}.sh`);
		let keep = false;
		try {
			const handle = await fs.open(snapshotPath, "wx", PRIVATE_FILE_MODE);
			try {
				await handle.chmod(PRIVATE_FILE_MODE);
			} finally {
				await handle.close();
			}

			const home = typeof options.home === "function" ? options.home() : options.home;
			const rcFile = getShellConfigFile(shell, home);
			const script = generateSnapshotScript(shell, snapshotPath, rcFile, await Bun.file(rcFile).exists());
			const snapshotEnv = sanitizeSnapshotEnv(env);
			const spawnEnv: Record<string, string> = {};
			for (const [key, value] of Object.entries(snapshotEnv)) {
				if (value !== undefined) spawnEnv[key] = value;
			}
			const child = Bun.spawn([shell, "-c", script], {
				env: spawnEnv,
				stdin: "ignore",
				stdout: "ignore",
				stderr: "ignore",
				timeout: SNAPSHOT_TIMEOUT_MS,
				killSignal: "SIGKILL",
			});

			await child.exited;
			if (child.exitCode === 0 && (await isTrustedSnapshot(snapshotPath))) {
				cachedPaths.set(shell, snapshotPath);
				keep = true;
				return snapshotPath;
			}
		} catch {
			// Snapshot creation failed, proceed without it.
		} finally {
			if (!keep) await fs.unlink(snapshotPath).catch(() => {});
		}
		return null;
	}

	return {
		async getOrCreateSnapshot(shell: string, env: Record<string, string | undefined>) {
			if (options.platform === "win32" || cleaningUp) return null;
			const currentLifecycle = lifecycle;
			const cached = cachedPaths.get(shell);
			if (cached && (await isTrustedSnapshot(cached))) {
				return cleaningUp || lifecycle !== currentLifecycle ? null : cached;
			}
			if (cleaningUp || lifecycle !== currentLifecycle) return null;
			if (cached) cachedPaths.delete(shell);
			const pending = inFlight.get(shell);
			if (pending) {
				const snapshot = await pending;
				return cleaningUp || lifecycle !== currentLifecycle ? null : snapshot;
			}
			const attempt = createSnapshot(shell, env);
			inFlight.set(shell, attempt);
			try {
				const snapshot = await attempt;
				return cleaningUp || lifecycle !== currentLifecycle ? null : snapshot;
			} finally {
				if (inFlight.get(shell) === attempt) inFlight.delete(shell);
			}
		},
		async cleanup() {
			if (cleanupAttempt) return await cleanupAttempt;
			const attempt = (async () => {
				cleaningUp = true;
				lifecycle += 1;
				try {
					const pending: Promise<unknown>[] = [...inFlight.values()];
					if (rootInitialization) pending.push(rootInitialization);
					await Promise.allSettled(pending);
					const root = rootPath ?? (await rootInitialization?.catch(() => null)) ?? null;
					cachedPaths.clear();
					inFlight.clear();
					rootPath = null;
					rootInitialization = null;
					if (root) await fs.rm(root, { recursive: true, force: true });
				} finally {
					cleaningUp = false;
				}
			})();
			cleanupAttempt = attempt;
			try {
				await attempt;
			} finally {
				if (cleanupAttempt === attempt) cleanupAttempt = null;
			}
		},
	};
}

/** @internal Test-only factory for isolated filesystem and platform assertions. */
export const createShellSnapshotCacheForTesting = createShellSnapshotCache;

const processSnapshotCache = createShellSnapshotCache({
	tempRoot: os.tmpdir(),
	home: () => os.homedir(),
	platform: process.platform,
});

/**
 * Create a shell snapshot, caching the result.
 * Returns the path to the snapshot file, or null if creation failed.
 */
export async function getOrCreateSnapshot(
	shell: string,
	env: Record<string, string | undefined>,
): Promise<string | null> {
	return await processSnapshotCache.getOrCreateSnapshot(shell, env);
}

postmortem.register("shell-snapshot", () => processSnapshotCache.cleanup());