import { randomBytes, randomUUID } from "node:crypto";
import { appendFile, mkdir, writeFile } from "node:fs/promises";
import { homedir } from "node:os";
import { dirname, join } from "node:path";
import { fileURLToPath } from "node:url";
import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { loadTraceConfig, writeGlobalTraceEnabled } from "./config.ts";
import { enforceTraceRetention } from "./retention.ts";
import { TraceRuntime, type TraceEventSink } from "./runtime.ts";
import { TraceWriter } from "./trace-writer.ts";
import { ViewerManager } from "./viewer-manager.ts";

function isAssistantMessage(message: unknown): boolean {
	return typeof message === "object" && message !== null && "role" in message && message.role === "assistant";
}

function notify(ctx: ExtensionContext, message: string, type: "info" | "warning" = "warning"): void {
	if (!ctx.hasUI) return;
	try { ctx.ui.notify(message, type); } catch { /* UI notification is best effort. */ }
}

function observe(ctx: ExtensionContext, action: () => void): void {
	try { action(); } catch (error) {
		notify(ctx, `Trace observer failed: ${error instanceof Error ? error.message : String(error)}`);
	}
}

async function createRuntime(ctx: ExtensionContext, reason: string): Promise<TraceRuntime> {
	const { config, warnings } = loadTraceConfig({ homeDir: homedir(), cwd: ctx.cwd });
	for (const warning of warnings) notify(ctx, warning);
	const sessionId = ctx.sessionManager.getSessionId();
	const sessionDir = ctx.sessionManager.getSessionDir();
	const sessionFile = ctx.sessionManager.getSessionFile();
	const runId = randomUUID();
	const now = new Date();
	const runName = `${now.toISOString().replaceAll(":", "-")}_${runId}`;
	const runPath = join(sessionDir, "traces", sessionId, "runs", runName);
	const tracePath = sessionFile === undefined ? undefined : join(runPath, "events.jsonl");
	const warning = (message: string): void => {
		notify(ctx, message);
		if (tracePath !== undefined) {
			void mkdir(runPath, { recursive: true })
				.then(() => appendFile(join(runPath, "viewer.log"), `${new Date().toISOString()} ${message}\n`, "utf8"))
				.catch(() => undefined);
		}
	};

	let persistence: TraceEventSink | undefined;
	if (tracePath !== undefined && config.persistence) {
		const writer = new TraceWriter(tracePath, (error) => warning(`Trace storage disabled: ${error instanceof Error ? error.message : String(error)}`));
		persistence = {
			write: (_event, line) => { writer.enqueue(line); },
			flush: (timeoutMs) => writer.flush(timeoutMs),
		};
	}

	const viewerManager = new ViewerManager({
		bootstrapPath: fileURLToPath(new URL("../viewer/bootstrap.mjs", import.meta.url)),
		tracePath,
		sessionId,
		token: randomBytes(24).toString("hex"),
		port: config.viewerPort,
		maxQueueBytes: config.maxQueueBytes,
		autoOpen: config.autoOpen,
		startupTimeoutMs: 5_000,
		onStatus: (status) => {
			if (status === "failed") warning("Trace Viewer unavailable; Pi and trace persistence remain active");
		},
		onWarning: warning,
	});
	const viewer: TraceEventSink = {
		write: (event, line) => { viewerManager.enqueue(event, line); },
		start: () => viewerManager.start(),
		shutdown: (timeoutMs) => viewerManager.shutdown(timeoutMs),
		open: () => viewerManager.open(),
	};
	const runtime = new TraceRuntime({
		sessionId,
		runId,
		config,
		clock: { now: () => new Date() },
		ids: { next: () => randomUUID() },
		persistence,
		viewer,
		onWarning: warning,
	});
	runtime.start(reason);

	if (tracePath !== undefined) {
		void (async () => {
			try {
				const metadataPath = join(sessionDir, "traces", sessionId, "metadata.json");
				await mkdir(dirname(metadataPath), { recursive: true });
				await writeFile(metadataPath, `${JSON.stringify({ schemaVersion: 1, sessionId, sessionFile, projectPath: ctx.cwd, createdAt: now.toISOString(), lastActivityAt: now.toISOString(), model: ctx.model ? { provider: ctx.model.provider, id: ctx.model.id } : undefined }, null, 2)}\n`);
				const result = await enforceTraceRetention({ sessionDir, activeRunPath: runPath, retentionDays: config.retentionDays, maxProjectBytes: config.maxProjectBytes, now });
				for (const retentionWarning of result.warnings) warning(retentionWarning);
			} catch (error) {
				warning(`Trace maintenance failed: ${error instanceof Error ? error.message : String(error)}`);
			}
		})();
	}
	return runtime;
}

export default function traceViewerExtension(pi: ExtensionAPI): void {
	let runtime: TraceRuntime | undefined;

	pi.on("session_start", async (event, ctx) => {
		try {
			if (runtime !== undefined) await runtime.shutdown("session replaced", 250);
			runtime = await createRuntime(ctx, event.reason);
		} catch (error) {
			runtime = undefined;
			notify(ctx, `Trace extension failed to start: ${error instanceof Error ? error.message : String(error)}`);
		}
	});
	pi.on("before_agent_start", (event, ctx) => {
		observe(ctx, () => runtime?.startAgentInput(event.prompt, event.images));
	});
	pi.on("message_start", (event, ctx) => observe(ctx, () => runtime?.observeMessage(event.message)));
	pi.on("agent_end", (_event, ctx) => observe(ctx, () => runtime?.endAgent()));
	pi.on("turn_start", (event, ctx) => observe(ctx, () => runtime?.startTurn(event.turnIndex, event.timestamp)));
	pi.on("turn_end", (event, ctx) => observe(ctx, () => runtime?.endTurn(event.turnIndex, event.message, event.toolResults)));
	pi.on("before_provider_request", (event, ctx) => { observe(ctx, () => runtime?.startLlm(event.payload)); });
	pi.on("message_end", (event, ctx) => {
		if (isAssistantMessage(event.message)) observe(ctx, () => runtime?.finishLlm(event.message));
	});
	pi.on("tool_execution_start", (event, ctx) => observe(ctx, () => runtime?.startTool(event.toolCallId, event.toolName, event.args)));
	pi.on("tool_execution_end", (event, ctx) => observe(ctx, () => runtime?.finishTool(event.toolCallId, event.toolName, event.result, event.isError)));
	pi.on("session_shutdown", async (event, ctx) => {
		const current = runtime;
		runtime = undefined;
		try { await current?.shutdown(event.reason, 250); } catch (error) {
			notify(ctx, `Trace shutdown failed: ${error instanceof Error ? error.message : String(error)}`);
		}
	});

	const handleTraceCommand = async (args: string, ctx: ExtensionCommandContext): Promise<void> => {
		const command = args.trim() || "status";
		if (command === "open") runtime?.open();
		else if (command === "on" || command === "off") {
			const enabled = command === "on";
			runtime?.setEnabled(enabled);
			try {
				await writeGlobalTraceEnabled({ homeDir: homedir(), enabled });
			} catch (error) {
				notify(ctx, `Trace observer changed for this session, but the global setting was not saved: ${error instanceof Error ? error.message : String(error)}`);
			}
		}
		else if (command !== "status") {
			notify(ctx, "Usage: /trace-viewer [status|open|on|off]", "info");
			return;
		}
		notify(ctx, `Trace observer: ${runtime?.isEnabled() ? "on" : "off"}`, "info");
	};
	const commandOptions = {
		description: "Control the local read-only trace observer: status, open, on, off",
		handler: handleTraceCommand,
	};
	pi.registerCommand("trace-viewer", commandOptions);
}