/**
 * The sandbox_run tool: execute code in a fresh disposable Docker sandbox.
 *
 * Lifecycle per invocation:
 *   1. Validate parameters (mode, timeout, artifacts).
 *   2. Resolve network policy (fail closed without UI; one-run human approval otherwise).
 *   3. Ensure the sandbox image exists; error with remediation if missing.
 *   4. Create a fresh hardened container (read-only /workspace, no secrets).
 *   5. Run the privileged setup phase as root (if any).
 *   6. Run the main command as the unprivileged sandbox user.
 *   7. Export requested artifacts from /output.
 *   8. Always destroy the container (finally), including failure/cancel/timeout.
 */

import { mkdtemp, rm, stat, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
	DEFAULT_MAX_BYTES,
	DEFAULT_MAX_LINES,
	formatSize,
	truncateTail,
	type TruncationResult,
} from "@earendil-works/pi-coding-agent";
import {
	ArtifactBudget,
	ArtifactError,
	guestArtifactPath,
	normalizeArtifactPath,
} from "./artifacts.ts";
import { createSandboxContainer, execInContainer, extractArtifact, imageExists, removeContainer } from "./docker.ts";
import { buildApprovalPrompt, decideNetwork } from "./permissions.ts";
import {
	ParamsError,
	normalizeArtifactList,
	resolveTimeoutSeconds,
	sandboxRunParamsSchema,
	validateExecutionMode,
	wantsNetwork,
	type SandboxRunParams,
} from "./params.ts";
import { DEFAULT_TIMEOUT_SECONDS, FS, IMAGE_TAG, MAX_ARTIFACT_BYTES, type ExportedArtifact, type SandboxRunDetails } from "./types.ts";

/** Capture + truncation helper for a single output stream. */
interface CapturedStream {
	text: string;
	truncated: TruncationResult | undefined;
	fullPath: string | undefined;
}

/**
 * Accumulate a stream, truncate to Pi's display limits, and persist the full
 * output to a temp file when truncated (so the model can read it).
 */
async function captureStream(
	messages: string[],
	prefix: string,
): Promise<CapturedStream> {
	const full = messages.join("");
	const truncation = truncateTail(full, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES });
	let fullPath: string | undefined;
	if (truncation.truncated && full) {
		const dir = await mkdtemp(join(tmpdir(), "pi-sandbox-"));
		fullPath = join(dir, `${prefix}.txt`);
		await writeFile(fullPath, full, "utf8");
	}
	return { text: truncation.content, truncated: truncation.truncated ? truncation : undefined, fullPath };
}

/** Print-mode / JSON-mode notification helper (no-op when there is no UI). */

/** Orchestrate a single sandbox run. */
export async function runSandbox(
	params: SandboxRunParams,
	ctx: ExtensionContext,
	signal: AbortSignal | undefined,
): Promise<{ content: { type: "text"; text: string }[]; details: SandboxRunDetails }> {
	const startedAt = Date.now();
	const mode = validateExecutionMode(params);
	const timeoutSeconds = resolveTimeoutSeconds(params.timeout_seconds);
	const artifactRequests = normalizeArtifactList(params.artifacts);
	const networkRequested = wantsNetwork(params);

	// ---- Network policy ---------------------------------------------------
	let networkApproved = false;
	if (networkRequested) {
		if (!ctx.hasUI) {
			const msg =
				"Network access was requested but there is no interactive UI to approve it. " +
				"Refusing to start a network-enabled container.";
			return {
				content: [{ type: "text", text: msg }],
				details: {
					exitCode: null,
					timedOut: false,
					cancelled: false,
					durationMs: Date.now() - startedAt,
					containerId: "",
					networkEnabled: false,
					setupRan: false,
					stdoutTruncated: false,
					stderrTruncated: false,
					artifacts: [],
				},
			};
		}
		const prompt = buildApprovalPrompt({
			command: params.command,
			executable: params.executable,
			args: params.args,
			setup: params.setup,
			projectPath: ctx.cwd,
		});
		const choice = await ctx.ui.select(prompt, ["Deny", "Allow once (this run only)"]);
		const decision = decideNetwork(true, true, choice === "Allow once (this run only)");
		if (!decision.allowed) {
			return {
				content: [{ type: "text", text: `Network access denied. ${decision.reason}` }],
				details: {
					exitCode: null,
					timedOut: false,
					cancelled: false,
					durationMs: Date.now() - startedAt,
					containerId: "",
					networkEnabled: false,
					setupRan: false,
					stdoutTruncated: false,
					stderrTruncated: false,
					artifacts: [],
				},
			};
		}
		networkApproved = true;
	}

	// ---- Image availability ----------------------------------------------
	const image = IMAGE_TAG;
	if (!(await imageExists(image))) {
		const msg =
			`Sandbox image ${image} is missing. ` +
			`Run /sandbox-build to build it, then retry.`;
		ctx.ui.notify(msg, "error");
		throw new ParamsError(msg);
	}

	// ---- Create + run container ------------------------------------------
	let containerId = "";
	let setupRan = false;
	const stdoutMessages: string[] = [];
	const stderrMessages: string[] = [];
	const artifacts: ExportedArtifact[] = [];
	const artifactBudget = new ArtifactBudget();

	const appendStdout = (chunk: string) => stdoutMessages.push(chunk);
	const appendStderr = (chunk: string) => stderrMessages.push(chunk);

	try {
		const container = await createSandboxContainer({
			image,
			projectPath: ctx.cwd,
			network: networkApproved,
		});
		containerId = container.id;

		// Setup phase (root, privileged) — e.g. package installation.
		if (params.setup && params.setup.length > 0) {
			for (const cmd of params.setup) {
				const res = await execInContainer(
					container,
					{ mode: "shell", command: cmd },
					{ cwd: FS.workspace, asUser: false, signal, timeoutSeconds: timeoutSeconds },
				);
				setupRan = true;
				if (res.exitCode !== 0) {
					throw new Error(`Setup command failed (exit ${res.exitCode}): ${cmd}\n${res.stdout}\n${res.stderr}`);
				}
			}
		}

		// Main command as the unprivileged sandbox user.
		const run =
			mode === "shell"
				? { mode: "shell" as const, command: params.command as string }
				: { mode: "argv" as const, executable: params.executable as string, args: params.args ?? [] };

		const execResult = await execInContainer(container, run, {
			cwd: FS.workspace,
			asUser: true,
			signal,
			timeoutSeconds,
			onStdout: appendStdout,
			onStderr: appendStderr,
		});

		// Export artifacts from /output.
		for (const requested of artifactRequests) {
			let normalized: string;
			try {
				normalized = normalizeArtifactPath(requested);
			} catch (err) {
				if (err instanceof ArtifactError) {
					artifacts.push({
						requested,
						guestPath: "",
						hostPath: "",
						bytes: 0,
					});
					stderrMessages.push(`[artifact rejected: ${err.message}]`);
					continue;
				}
				throw err;
			}
			const guestPath = guestArtifactPath(normalized);
			const hostDir = await mkdtemp(join(tmpdir(), "pi-sandbox-art-"));
			const hostPath = join(hostDir, normalized.replace(/\//g, "__"));
			try {
				await extractArtifact(container, guestPath, hostPath);
				const size = (await stat(hostPath).catch(() => null))?.size ?? 0;
				if (!artifactBudget.tryConsume(size)) {
					stderrMessages.push(`[artifact skipped: ${normalized} exceeds aggregate budget of ${formatSize(MAX_ARTIFACT_BYTES)}]`);
					await rm(hostPath, { recursive: true, force: true });
					continue;
				}
				artifacts.push({ requested: normalized, guestPath, hostPath, bytes: size });
			} catch {
				stderrMessages.push(`[artifact not found: ${normalized}]`);
			}
		}

		// ---- Build result -------------------------------------------------
		const stdoutCaptured = await captureStream(stdoutMessages, "stdout");
		const stderrCaptured = await captureStream(stderrMessages, "stderr");

		const body: string[] = [];
		if (stdoutCaptured.text) body.push(stdoutCaptured.text);
		if (stderrCaptured.text) body.push(`[stderr]\n${stderrCaptured.text}`);

		const meta: string[] = [];
		meta.push(`exit code: ${execResult.exitCode ?? "n/a"}`);
		meta.push(`duration: ${((Date.now() - startedAt) / 1000).toFixed(1)}s`);
		meta.push(`network: ${networkApproved ? "enabled (approved)" : "disabled"}`);
		if (setupRan) meta.push("setup: ran as root");
		if (stdoutCaptured.truncated) meta.push("stdout truncated");
		if (stderrCaptured.truncated) meta.push("stderr truncated");
		if (execResult.timedOut) meta.push(`timed out after ${timeoutSeconds}s`);
		if (artifacts.length > 0) {
			meta.push(`artifacts: ${artifacts.map((a) => `${a.requested} -> ${a.hostPath}`).join(", ")}`);
		}
		body.push(`[${meta.join(". ")}]`);

		const details: SandboxRunDetails = {
			exitCode: execResult.exitCode,
			timedOut: execResult.timedOut,
			cancelled: false,
			durationMs: Date.now() - startedAt,
			containerId,
			networkEnabled: networkApproved,
			setupRan,
			stdoutTruncated: !!stdoutCaptured.truncated,
			stderrTruncated: !!stderrCaptured.truncated,
			stdoutPath: stdoutCaptured.fullPath,
			stderrPath: stderrCaptured.fullPath,
			artifacts,
		};

		return {
			content: [{ type: "text", text: body.join("\n\n") }],
			details,
		};
	} catch (err) {
		const cancelled = signal?.aborted ?? false;
		// A timeout surfaced as an error (e.g. a setup command timing out) should be
		// reported as a timeout, not a generic failure.
		const timedOut = err instanceof Error && /timed? ?out/i.test(err.message);
		const stdoutCaptured = await captureStream(stdoutMessages, "stdout");
		const stderrCaptured = await captureStream(stderrMessages, "stderr");

		const body: string[] = [];
		if (stdoutCaptured.text) body.push(stdoutCaptured.text);
		if (stderrCaptured.text) body.push(`[stderr]\n${stderrCaptured.text}`);
		body.push(`[error: ${err instanceof Error ? err.message : String(err)}]`);

		const details: SandboxRunDetails = {
			exitCode: null,
			timedOut,
			cancelled,
			durationMs: Date.now() - startedAt,
			containerId,
			networkEnabled: networkApproved,
			setupRan,
			stdoutTruncated: !!stdoutCaptured.truncated,
			stderrTruncated: !!stderrCaptured.truncated,
			stdoutPath: stdoutCaptured.fullPath,
			stderrPath: stderrCaptured.fullPath,
			artifacts,
		};

		return {
			content: [{ type: "text", text: body.join("\n\n") }],
			details,
		};
	} finally {
		if (containerId) {
			await removeContainer(containerId);
		}
	}
}

/** Register the sandbox_run tool on the given ExtensionAPI. */
export function registerSandboxTool(pi: ExtensionAPI): void {
	pi.registerTool({
		name: "sandbox_run",
		label: "sandbox_run",
		description: `Execute code in a fresh, disposable, offline Debian sandbox container. ` +
			`Use for building, testing, compiling, or running code that must not touch the host. ` +
			`The host project is mounted read-only at /workspace; write exportable results under /output. ` +
			`Networking is OFF by default; setting network:true requires human approval for that single run. ` +
			`Provide exactly one of command (shell) or executable+args (no shell).`,
		parameters: sandboxRunParamsSchema,
		async execute(_toolCallId, params, signal, _onUpdate, ctx) {
			return runSandbox(params, ctx, signal);
		},
		promptSnippet: "Run code in a fresh disposable offline Debian sandbox (read-only /workspace, /output for artifacts)",
		promptGuidelines: [
			"Use sandbox_run for building, testing, or running code that should not touch the host.",
			`Timeout defaults to ${DEFAULT_TIMEOUT_SECONDS}s; pass timeout_seconds to change it.`,
			"Networking is off by default and requires human approval per run.",
		],
	});
}