/**
 * Persistent bash shell — the pi counterpart of DSH's
 * `@deepseek-ai/dsh-tool-bash-persistent`.
 *
 * DSH drives a PTY; pi has no PTY dependency, so this drives one long-lived
 * `bash` process over pipes with the same marker protocol DSH uses:
 *
 *   printf '%s\n' START; eval -- $'cmd'; status=$?; printf '%s%s\n' END "$status"
 *
 * stderr is merged into stdout inside the shell (`exec 2>&1`) so the model
 * sees interleaved output like it would on a terminal. Each command's stdin
 * is `/dev/null` so a stray `cat`/`read` cannot swallow the next command.
 *
 * Semantics kept from DSH:
 *  - state (cwd, env, functions) persists across calls;
 *  - one command at a time per shell (serialized queue);
 *  - timeout → partial output + the shell is reset;
 *  - abort → shell reset, error propagates;
 *  - shell death → partial output + exit status + reset notice;
 *  - non-zero exit → `[exit code: N]` marker appended;
 *  - output truncated to `maxOutputChars` with the DSH clip note.
 */

import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
import { randomUUID } from "node:crypto";
import { maybeTruncate } from "./schemas.ts";

export const SHELL_RESET_MESSAGE =
	"The persistent bash shell was reset; the next bash call starts from the workspace with a fresh current directory and environment.";

export interface PersistentShellOptions {
	cwd: string;
	bashPath?: string;
	timeoutMs?: number;
	maxOutputChars?: number;
	env?: NodeJS.ProcessEnv;
}

export interface CommandOutcome {
	/** Text handed to the model. */
	text: string;
	exitCode?: number;
	timedOut: boolean;
	shellExited: boolean;
}

function quoteForBash(value: string): string {
	return `$'${value.replaceAll("\\", "\\\\").replaceAll("'", "\\'").replaceAll("\r", "\\r").replaceAll("\n", "\\n")}'`;
}

function markers() {
	const nonce = randomUUID();
	return {
		start: `__PI_PERSISTENT_BASH_START_${nonce}__`,
		end: `__PI_PERSISTENT_BASH_END_${nonce}:`,
	};
}

export function wrapCommand(command: string, marker: { start: string; end: string }): string {
	return (
		`printf '%s\\n' ${quoteForBash(marker.start)}; ` +
		`eval -- ${quoteForBash(command)} < /dev/null; ` +
		`__pi_persistent_bash_status=$?; ` +
		`printf '%s%s\\n' ${quoteForBash(marker.end)} "$__pi_persistent_bash_status"`
	);
}

function appendStatusMarker(content: string, marker: string | undefined): string {
	if (marker === undefined) return content;
	return content.length === 0 ? marker : `${content}\n${marker}`;
}

/** Extract the completed command's output from the shell's raw stream. */
export function parseCompleted(
	buffer: string,
	marker: { start: string; end: string },
): { text: string; exitCode: number } | undefined {
	const end = buffer.lastIndexOf(marker.end);
	if (end < 0) return undefined;
	const statusMatch = /^(\d+)\r?\n/.exec(buffer.slice(end + marker.end.length));
	if (!statusMatch) return undefined;
	const startAt = buffer.lastIndexOf(marker.start, end);
	const start = startAt < 0 ? 0 : startAt + marker.start.length;
	const text = buffer.slice(start, end).replace(/^\r?\n/, "").replace(/\r?\n$/, "");
	return { text, exitCode: Number(statusMatch[1]) };
}

/** Best-effort partial output when the command did not complete. */
export function parsePartial(buffer: string, marker: { start: string; end: string }): string {
	const startAt = buffer.lastIndexOf(marker.start);
	const after = startAt < 0 ? buffer : buffer.slice(startAt + marker.start.length);
	return after.replace(/^\r?\n/, "").replace(/\r?\n$/, "");
}

export class PersistentShell {
	private child: ChildProcessWithoutNullStreams | undefined;
	private buffer = "";
	private exited: { code: number | null; signal: NodeJS.Signals | null } | undefined;
	private queue: Promise<unknown> = Promise.resolve();
	private readonly opts: Required<Omit<PersistentShellOptions, "env">> & { env?: NodeJS.ProcessEnv };

	constructor(options: PersistentShellOptions) {
		this.opts = {
			cwd: options.cwd,
			bashPath: options.bashPath ?? "bash",
			timeoutMs: options.timeoutMs ?? 300_000,
			maxOutputChars: options.maxOutputChars ?? 16_000,
			env: options.env,
		};
	}

	get alive(): boolean {
		return this.child !== undefined && this.exited === undefined;
	}

	/** Run one command; calls are serialized per shell instance. */
	run(command: string, signal?: AbortSignal): Promise<CommandOutcome> {
		const task = () => this.execute(command, signal);
		const run = this.queue.then(task, task);
		this.queue = run.then(
			() => undefined,
			() => undefined,
		);
		return run;
	}

	/** Kill the shell; the next call spawns a fresh one. */
	async reset(): Promise<void> {
		const child = this.child;
		this.child = undefined;
		this.buffer = "";
		this.exited = undefined;
		if (child && child.exitCode === null && child.signalCode === null) {
			child.kill("SIGKILL");
		}
	}

	async dispose(): Promise<void> {
		await this.reset();
	}

	private ensureShell(): ChildProcessWithoutNullStreams {
		if (this.child && this.exited === undefined) return this.child;
		this.buffer = "";
		this.exited = undefined;
		const child = spawn(this.opts.bashPath, ["--noprofile", "--norc"], {
			cwd: this.opts.cwd,
			env: { ...(this.opts.env ?? process.env), TERM: "dumb", PS1: "", PROMPT_COMMAND: "" },
			stdio: ["pipe", "pipe", "pipe"],
		});
		child.stdout.setEncoding("utf8");
		child.stderr.setEncoding("utf8");
		child.stdout.on("data", (chunk: string) => {
			this.buffer += chunk;
		});
		// Before `exec 2>&1` takes effect (and if it ever fails), keep stderr visible.
		child.stderr.on("data", (chunk: string) => {
			this.buffer += chunk;
		});
		child.on("exit", (code, signal) => {
			if (this.child === child) this.exited = { code, signal };
		});
		child.on("error", () => {
			if (this.child === child) this.exited = { code: null, signal: null };
		});
		// Merge stderr into stdout inside the shell, disable history expansion so
		// `!` in commands is literal, and never let a command steal our stdin.
		child.stdin.write("exec 2>&1; set +H; set +m\n");
		this.child = child;
		return child;
	}

	private async execute(command: string, signal?: AbortSignal): Promise<CommandOutcome> {
		if (command.trim().length === 0) throw new Error("command must be a non-empty string");
		signal?.throwIfAborted();
		const child = this.ensureShell();
		const marker = markers();
		const wrapped = wrapCommand(command, marker);
		const startOffset = this.buffer.length;

		const startedAt = Date.now();

		try {
			child.stdin.write(`${wrapped}\n`);
		} catch (error) {
			await this.reset();
			throw error;
		}

		const view = () => this.buffer.slice(startOffset);
		try {
			for (;;) {
				const completed = parseCompleted(view(), marker);
				if (completed) {
					return {
						text: this.render(completed.text, completed.exitCode),
						exitCode: completed.exitCode,
						timedOut: false,
						shellExited: false,
					};
				}
				if (this.exited !== undefined || this.child !== child) {
					const partial = parsePartial(view(), marker);
					const status = this.exited;
					await this.reset();
					const shown = appendStatusMarker(
						maybeTruncate(partial, this.opts.maxOutputChars),
						status?.signal
							? `[shell killed by signal: ${status.signal}]`
							: status && status.code !== null
								? `[shell exited: code ${status.code}]`
								: "[shell exited]",
					);
					return {
						text: [shown, SHELL_RESET_MESSAGE].filter((p) => p.length > 0).join("\n"),
						timedOut: false,
						shellExited: true,
					};
				}
				if (signal?.aborted) {
					await this.reset();
					signal.throwIfAborted();
				}
				if (Date.now() - startedAt >= this.opts.timeoutMs) {
					const partial = maybeTruncate(parsePartial(view(), marker), this.opts.maxOutputChars);
					await this.reset();
					return {
						text: [
							`Your command timed out after ${Math.round(this.opts.timeoutMs / 1000)} seconds or experienced an OOM error. Below is partial output:`,
							partial,
							SHELL_RESET_MESSAGE,
						].join("\n"),
						timedOut: true,
						shellExited: false,
					};
				}
				await new Promise<void>((resolve) => setTimeout(resolve, 25));
			}
		} finally {
			// Drop consumed output so the buffer never grows unbounded.
			if (this.child === child) this.buffer = "";
		}
	}

	private render(text: string, exitCode: number): string {
		const rendered = maybeTruncate(text, this.opts.maxOutputChars);
		return appendStatusMarker(rendered, exitCode !== 0 ? `[exit code: ${exitCode}]` : undefined);
	}
}