import {
	AgentOs,
	type AgentOsOptions,
	type AgentOsSidecar,
	type RootSnapshotExport,
	type SoftwareInput,
} from "@rivet-dev/agentos-core";
import { coreutils } from "@agentos-software/common";
import claude from "@agentos-software/claude-code";
import pi from "@agentos-software/pi";
import { LLMock } from "@copilotkit/llmock";
import os from "node:os";

// Benchmark parameters. Keep batch sizes minimal for fast iteration.
export const ITERATIONS = 5;
export const WARMUP_ITERATIONS = 1;

export const ECHO_COMMAND = "echo hello";
export const EXPECTED_OUTPUT = "hello\n";
export const PI_BENCHMARK_PROMPT = "Reply with exactly: Hello from llmock";
export const PI_HEADLESS_BLOCKER_REFERENCE =
	"packages/core/tests/pi-headless.test.ts";
export const PI_HEADLESS_BLOCKER_REASON =
	'Standalone `spawn("pi", ...)` is not exposed on the native sidecar PATH; use `openSession({ sessionId: "main", agent: "pi-cli" })` to benchmark the native PI CLI RPC path tracked in packages/core/tests/pi-headless.test.ts.';
// ── Shared bench sidecar + cold-run snapshot ───────────────────────
//
// Benchmarks create the sidecar ONCE up front and lease every VM from it,
// rather than letting each `AgentOs.create()` stand up its own sidecar. A
// single cold-run VM is then created and its root filesystem snapshotted so
// the subsequent measured VMs boot from a warm snapshot instead of paying the
// bootstrap cost on every iteration. Both are wired through `benchCreateOptions`
// so every VM-creation helper picks them up automatically.

let _benchSidecar: AgentOsSidecar | undefined;
let _benchSnapshot: RootSnapshotExport | undefined;

/** Create the shared bench sidecar. Call once before creating any VM. */
export async function startBenchSidecar(): Promise<AgentOsSidecar> {
	if (!_benchSidecar) {
		_benchSidecar = await AgentOs.createSidecar();
	}
	return _benchSidecar;
}

/** Dispose the shared bench sidecar and clear the cold-run snapshot. */
export async function stopBenchSidecar(): Promise<void> {
	_benchSnapshot = undefined;
	if (_benchSidecar) {
		const sidecar = _benchSidecar;
		_benchSidecar = undefined;
		await sidecar.dispose();
	}
}

/** Record the cold-run root snapshot reused by subsequent measured VMs. */
export function setBenchRootSnapshot(snapshot: RootSnapshotExport): void {
	_benchSnapshot = snapshot;
}

/** Clear the cold-run root snapshot (e.g. when switching workloads). */
export function clearBenchRootSnapshot(): void {
	_benchSnapshot = undefined;
}

/**
 * Overlay the shared bench sidecar (and cold-run snapshot, when set) onto a
 * VM-creation options object. The snapshot only applies when the caller did
 * not request its own `rootFilesystem`.
 */
export function benchCreateOptions(options: AgentOsOptions = {}): AgentOsOptions {
	const overlay: AgentOsOptions = { ...options };
	if (_benchSidecar) {
		overlay.sidecar = { kind: "explicit", handle: _benchSidecar };
	}
	if (_benchSnapshot && overlay.rootFilesystem === undefined) {
		overlay.rootFilesystem = { type: "overlay", lowers: [_benchSnapshot] };
	}
	return overlay;
}

// ── Shared mock LLM server ─────────────────────────────────────────

let _llmock: LLMock | undefined;
let _llmockUrl: string | undefined;
let _llmockPort: number | undefined;

/** Start a shared llmock server (idempotent). */
export async function ensureLlmock(): Promise<{
	url: string;
	port: number;
}> {
	if (_llmock) return { url: _llmockUrl!, port: _llmockPort! };
	_llmock = new LLMock({ port: 0, logLevel: "silent" });
	_llmock.addFixtures([
		{
			match: { predicate: () => true },
			response: { content: "Hello from llmock" },
		},
	]);
	_llmockUrl = await _llmock.start();
	_llmockPort = Number(new URL(_llmockUrl).port);
	return { url: _llmockUrl, port: _llmockPort };
}

/** Stop the shared llmock server. */
export async function stopLlmock(): Promise<void> {
	if (_llmock) {
		await _llmock.stop();
		_llmock = undefined;
		_llmockUrl = undefined;
		_llmockPort = undefined;
	}
}

export function getLlmockRequestCount(): number {
	return _llmock?.getRequests().length ?? 0;
}

// ── Workload abstraction ────────────────────────────────────────────

export interface WorkloadObservation {
	promptCompleted?: boolean;
	providerRequestCount?: number;
	sessionUpdateCount?: number;
	textEventCount?: number;
	finalText?: string | null;
	stopReason?: string;
	workloadPath?: string;
	substituteReason?: string;
	blockerReference?: string;
	blockerReason?: string;
}

/** A workload describes how to create a VM and start a long-running process for memory measurement. */
export interface Workload {
	name: string;
	description: string;
	createVm: () => Promise<AgentOs>;
	/** Start a long-running process so the Worker thread stays alive. */
	start: (vm: AgentOs) => Promise<WorkloadObservation | void> | WorkloadObservation | void;
	/** Verify the expected processes are running. Throws if not. */
	verify: (vm: AgentOs) => void;
	/** Time to wait after start for the process to fully initialize. */
	settleMs: number;
}

function makeAgentSessionWorkload(opts: {
	agentId: string;
	description: string;
	software: SoftwareInput[];
	processMarker: string;
}): Workload {
	return {
		name: `${opts.agentId}-session`,
		description: opts.description,
		createVm: async () => {
			const { port } = await ensureLlmock();
			return AgentOs.create(
				benchCreateOptions({
					software: opts.software,
					loopbackExemptPorts: [port],
				}),
			);
		},
		start: async (vm) => {
			const { url } = await ensureLlmock();
			const sessionId = "main";
			await vm.openSession({
				sessionId,
				agent: opts.agentId,
				env: {
					ANTHROPIC_API_KEY: "bench-key",
					ANTHROPIC_BASE_URL: url,
				},
			});
		},
		verify: (vm) => {
			const procs = vm.listProcesses();
			const running = procs.filter((p) => p.running);
			const hasAgent = running.some(
				(p) =>
					p.command === "node" &&
					p.args.some((a) => a.includes(opts.processMarker)),
			);
			if (!hasAgent) {
				throw new Error(
					`Expected running ${opts.processMarker} process, got: ${JSON.stringify(running.map((p) => ({ cmd: p.command, args: p.args })))}`,
				);
			}
		},
		settleMs: 2000,
	};
}

function getTextEventPayload(
	event: unknown,
): { text?: string; type?: string } | undefined {
	if (!event || typeof event !== "object") {
		return undefined;
	}
	const update = (event as { update?: unknown }).update;
	if (!update || typeof update !== "object") {
		return undefined;
	}
	const content = (update as { content?: unknown }).content;
	if (!content || typeof content !== "object") {
		return undefined;
	}
	return content as { text?: string; type?: string };
}

function makeAgentPromptWorkload(opts: {
	agentId: string;
	description: string;
	software: SoftwareInput[];
	processMarker: string;
	prompt: string;
}): Workload {
	return {
		name: `${opts.agentId}-prompt-turn`,
		description: opts.description,
		createVm: async () => {
			const { port } = await ensureLlmock();
			return AgentOs.create(
				benchCreateOptions({
					loopbackExemptPorts: [port],
					software: opts.software,
				}),
			);
		},
		start: async (vm) => {
			const { url } = await ensureLlmock();
			const sessionId = "main";
			await vm.openSession({ sessionId,
				agent: opts.agentId,
				env: {
					ANTHROPIC_API_KEY: "bench-key",
					ANTHROPIC_BASE_URL: url,
				},
			});

			const events: unknown[] = [];
			const unsubscribe = vm.onSessionEvent(sessionId, (event) => {
				events.push(event);
			});
			const requestCountBefore = getLlmockRequestCount();

			try {
				const { text, stopReason } = await vm.prompt({
					sessionId,
					content: [{ type: "text", text: opts.prompt }],
				});
				const textEvents = events
					.map(getTextEventPayload)
					.filter((event) => event?.type === "text");
				const finalText = textEvents.at(-1)?.text ?? text;
				const providerRequestCount =
					getLlmockRequestCount() - requestCountBefore;

				return {
					promptCompleted: true,
					providerRequestCount,
					sessionUpdateCount: events.length,
					textEventCount: textEvents.length,
					finalText,
					stopReason,
					workloadPath:
						'openSession({ sessionId: "main", agent: "pi-cli" }) + vm.prompt(...) via pi-acp -> PI CLI --mode rpc',
					blockerReference: PI_HEADLESS_BLOCKER_REFERENCE,
					blockerReason: PI_HEADLESS_BLOCKER_REASON,
				} satisfies WorkloadObservation;
			} finally {
				unsubscribe();
			}
		},
		verify: (vm) => {
			const procs = vm.listProcesses();
			const running = procs.filter((p) => p.running);
			const hasAgent = running.some(
				(p) =>
					p.command === "node" &&
					p.args.some((a) => a.includes(opts.processMarker)),
			);
			if (!hasAgent) {
				throw new Error(
					`Expected running ${opts.processMarker} process, got: ${JSON.stringify(running.map((p) => ({ cmd: p.command, args: p.args })))}`,
				);
			}
		},
		settleMs: 2000,
	};
}

export const WORKLOADS: Record<string, Workload> = {
	sleep: {
		name: "sleep",
		description: "Minimal VM with idle Node.js process (setTimeout keepalive)",
		createVm: () => AgentOs.create(benchCreateOptions()),
		start: (vm) => {
			vm.spawn("node", ["-e", "setTimeout(() => {}, 999999999)"], {
				streamStdin: true,
			});
		},
		verify: (vm) => {
			const procs = vm.listProcesses();
			const running = procs.filter((p) => p.running);
			const hasNode = running.some((p) => p.command === "node");
			if (!hasNode) {
				throw new Error(
					`Expected running 'node' process, got: ${JSON.stringify(running.map((p) => p.command))}`,
				);
			}
		},
		settleMs: 2000,
	},
	"pi-session": makeAgentSessionWorkload({
		agentId: "pi",
		description: "VM with PI agent session via openSession",
		software: [pi],
		processMarker: "agentos-pi",
	}),
	"pi-prompt-turn": makeAgentPromptWorkload({
		agentId: "pi-cli",
		description:
			'Native PI CLI headless benchmark path via openSession({ sessionId: "main", agent: "pi-cli" }), which drives the real PI CLI through pi-acp RPC mode and records a full prompt turn.',
		software: [],
		processMarker: "pi-acp",
		prompt: PI_BENCHMARK_PROMPT,
	}),
	"claude-session": makeAgentSessionWorkload({
		agentId: "claude",
		description: "VM with Claude agent session via openSession",
		software: [claude],
		processMarker: "agentos-claude",
	}),
};

// ── VM creation helpers ─────────────────────────────────────────────

/**
 * Create a fresh AgentOS VM with only coreutils (WASM shell + echo).
 * This is the minimal setup needed to run shell commands.
 */
export async function createBenchVm(): Promise<AgentOs> {
	return AgentOs.create(
		benchCreateOptions({
			software: [coreutils],
		}),
	);
}

// ── Stats and formatting ────────────────────────────────────────────

export function percentile(sorted: number[], p: number): number {
	const idx = Math.ceil((p / 100) * sorted.length) - 1;
	return sorted[Math.max(0, idx)];
}

export function stats(samples: number[]) {
	const sorted = [...samples].sort((a, b) => a - b);
	const mean = samples.reduce((a, b) => a + b, 0) / samples.length;
	return {
		mean: round(mean),
		p50: round(percentile(sorted, 50)),
		p95: round(percentile(sorted, 95)),
		p99: round(percentile(sorted, 99)),
		min: round(sorted[0]),
		max: round(sorted[sorted.length - 1]),
	};
}

export function round(n: number, decimals = 2): number {
	const f = 10 ** decimals;
	return Math.round(n * f) / f;
}

export function getHardware() {
	const cpus = os.cpus();
	return {
		cpu: cpus[0]?.model ?? "unknown",
		cores: os.availableParallelism(),
		ram: `${round(os.totalmem() / 1024 ** 3, 1)} GB`,
		node: process.version,
		os: `${os.type()} ${os.release()}`,
		arch: os.arch(),
	};
}

export function forceGC() {
	if (global.gc) {
		global.gc();
	} else {
		console.error("WARNING: global.gc not available. Run with --expose-gc");
	}
}

export async function sleep(ms: number): Promise<void> {
	return new Promise((r) => setTimeout(r, ms));
}

export function formatBytes(bytes: number): string {
	if (Math.abs(bytes) < 1024) return `${bytes} B`;
	const mb = bytes / (1024 * 1024);
	return `${round(mb, 2)} MB`;
}

/** Print a table to stderr for human readability. */
export function printTable(
	headers: string[],
	rows: (string | number)[][],
): void {
	const widths = headers.map((h, i) =>
		Math.max(h.length, ...rows.map((r) => String(r[i]).length)),
	);
	const sep = widths.map((w) => "-".repeat(w)).join(" | ");
	const fmt = (row: (string | number)[]) =>
		row.map((c, i) => String(c).padStart(widths[i])).join(" | ");

	console.error("");
	console.error(fmt(headers));
	console.error(sep);
	for (const row of rows) {
		console.error(fmt(row));
	}
	console.error("");
}