#!/usr/bin/env node

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";

interface TranscriptRecord {
	time: number;
	order: number;
	session: string;
	role: "system" | "user" | "assistant" | "tool";
	source?: string;
	content: string;
}

interface Descriptor {
	name?: string;
	sessionId?: string;
	sessionFile?: string;
	createdAt?: string;
	appendSystemPrompt?: string;
}

function usage(): never {
	throw new Error(
		"usage: node --experimental-strip-types scripts/agi-compact-run.ts <run-directory> [output-file] [--prune]",
	);
}

function walkFiles(directory: string): string[] {
	const files: string[] = [];
	for (const entry of fs.readdirSync(directory, { withFileTypes: true })) {
		const target = path.join(directory, entry.name);
		if (entry.isDirectory()) files.push(...walkFiles(target));
		else if (entry.isFile()) files.push(target);
	}
	return files;
}

function parseJsonLines(file: string): unknown[] {
	const rows: unknown[] = [];
	for (const line of fs.readFileSync(file, "utf8").split("\n")) {
		if (!line.trim()) continue;
		try {
			rows.push(JSON.parse(line));
		} catch {
			// A compact transcript should preserve valid session content, not malformed runtime fragments.
		}
	}
	return rows;
}

function record(value: unknown): Record<string, unknown> | undefined {
	return typeof value === "object" && value !== null ? (value as Record<string, unknown>) : undefined;
}

function textContent(value: unknown): string {
	if (typeof value === "string") return value;
	if (!Array.isArray(value)) return value === undefined ? "" : JSON.stringify(value);
	const parts: string[] = [];
	for (const item of value) {
		const part = record(item);
		if (!part) continue;
		if (part.type === "text" && typeof part.text === "string") parts.push(part.text);
		else if (part.type === "image") parts.push("<image>");
		else parts.push(JSON.stringify(part));
	}
	return parts.join("\n");
}

function assistantContent(value: unknown): string {
	if (!Array.isArray(value)) return textContent(value);
	const parts: string[] = [];
	for (const item of value) {
		const part = record(item);
		if (!part) continue;
		if (part.type === "thinking" && typeof part.thinking === "string") {
			if (part.thinking) parts.push(`thinking:\n${part.thinking}`);
		} else if (part.type === "text" && typeof part.text === "string") {
			if (part.text) parts.push(part.text);
		} else if (part.type === "toolCall") {
			const name = typeof part.name === "string" ? part.name : "unknown";
			parts.push(`tool ${name} ${JSON.stringify(part.arguments ?? {})}`);
		} else if (part.type === "image") {
			parts.push("<image>");
		} else {
			const semantic = { ...part };
			delete semantic.id;
			delete semantic.thinkingSignature;
			parts.push(JSON.stringify(semantic));
		}
	}
	return parts.join("\n\n");
}

function timestamp(value: unknown, fallback: number): number {
	if (typeof value === "string") {
		const parsed = Date.parse(value);
		if (Number.isFinite(parsed)) return parsed;
	}
	if (typeof value === "number" && Number.isFinite(value)) return value;
	return fallback;
}

function sessionLabels(runDirectory: string, sessionFiles: string[]): Map<string, string> {
	const labels = new Map<string, string>();
	const descriptors = walkFiles(runDirectory).filter((file) => file.endsWith("/descriptor.json"));
	for (const file of descriptors) {
		let descriptor: Descriptor;
		try {
			descriptor = JSON.parse(fs.readFileSync(file, "utf8")) as Descriptor;
		} catch {
			continue;
		}
		const label = descriptor.name ? `worker:${descriptor.name}` : "worker";
		if (descriptor.sessionId) labels.set(descriptor.sessionId, label);
		if (descriptor.sessionFile) labels.set(path.basename(descriptor.sessionFile), label);
	}
	for (const file of sessionFiles) {
		const session = parseJsonLines(file).map(record).find((entry) => entry?.type === "session");
		const id = typeof session?.id === "string" ? session.id : path.basename(file, ".jsonl");
		if (!labels.has(id) && !labels.has(path.basename(file))) {
			labels.set(id, id.includes("root") ? "root" : id);
		}
	}
	return labels;
}

function compactRun(runDirectory: string): { transcript: string; recordCount: number; systemPromptCaptured: boolean } {
	const sessionRoot = path.join(runDirectory, "raw", "agent", "pi-config", "sessions");
	if (!fs.existsSync(sessionRoot)) throw new Error(`session directory not found: ${sessionRoot}`);
	const sessionFiles = walkFiles(sessionRoot).filter((file) => file.endsWith(".jsonl")).sort();
	if (sessionFiles.length === 0) throw new Error(`no session JSONL files found under ${sessionRoot}`);
	const labels = sessionLabels(runDirectory, sessionFiles);
	const records: TranscriptRecord[] = [];
	let order = 0;
	let systemPromptCaptured = false;

	for (const file of sessionFiles) {
		const rows = parseJsonLines(file);
		const header = rows.map(record).find((entry) => entry?.type === "session");
		const sessionId = typeof header?.id === "string" ? header.id : path.basename(file, ".jsonl");
		const session = labels.get(sessionId) ?? labels.get(path.basename(file)) ?? sessionId;
		let fallbackTime = timestamp(header?.timestamp, 0);
		let previousSystemPrompt: string | undefined;
		for (const raw of rows) {
			const entry = record(raw);
			if (!entry) continue;
			fallbackTime += 1;
			const time = timestamp(entry.timestamp, fallbackTime);
			if (entry.type === "custom" && entry.customType === "system-prompt") {
				const data = record(entry.data);
				const content = typeof data?.systemPrompt === "string" ? data.systemPrompt : "";
				if (content && content !== previousSystemPrompt) {
					records.push({ time, order: order++, session, role: "system", content });
					previousSystemPrompt = content;
					systemPromptCaptured = true;
				}
				continue;
			}
			if (entry.type === "custom_message") {
				const content = textContent(entry.content);
				if (content) {
					records.push({
						time,
						order: order++,
						session,
						role: "user",
						source: typeof entry.customType === "string" ? entry.customType : "injected",
						content,
					});
				}
				continue;
			}
			if (entry.type === "compaction" && typeof entry.summary === "string") {
				records.push({ time, order: order++, session, role: "user", source: "compaction", content: entry.summary });
				continue;
			}
			if (entry.type === "branch_summary" && typeof entry.summary === "string") {
				records.push({ time, order: order++, session, role: "user", source: "branch-summary", content: entry.summary });
				continue;
			}
			if (entry.type !== "message") continue;
			const message = record(entry.message);
			if (!message) continue;
			const messageTime = timestamp(message.timestamp, time);
			if (message.role === "user") {
				const content = textContent(message.content);
				if (content) records.push({ time: messageTime, order: order++, session, role: "user", content });
			} else if (message.role === "assistant") {
				const content = assistantContent(message.content);
				if (content) records.push({ time: messageTime, order: order++, session, role: "assistant", content });
			} else if (message.role === "toolResult") {
				const content = textContent(message.content);
				const toolName = typeof message.toolName === "string" ? message.toolName : "unknown";
				const source = message.isError === true ? `${toolName}:error` : toolName;
				if (content) records.push({ time: messageTime, order: order++, session, role: "tool", source, content });
			}
		}
	}

	if (!systemPromptCaptured) {
		const seenAppendPrompts = new Set<string>();
		for (const file of walkFiles(runDirectory).filter((candidate) => candidate.endsWith("/descriptor.json"))) {
			let descriptor: Descriptor;
			try {
				descriptor = JSON.parse(fs.readFileSync(file, "utf8")) as Descriptor;
			} catch {
				continue;
			}
			if (!descriptor.appendSystemPrompt) continue;
			const session = descriptor.name ? `worker:${descriptor.name}` : "worker";
			const dedupeKey = `${descriptor.sessionId ?? session}\0${descriptor.appendSystemPrompt}`;
			if (seenAppendPrompts.has(dedupeKey)) continue;
			seenAppendPrompts.add(dedupeKey);
			records.push({
				time: timestamp(descriptor.createdAt, 0),
				order: order++,
				session,
				role: "system",
				source: "append-only; full prompt unavailable",
				content: descriptor.appendSystemPrompt,
			});
		}
	}

	records.sort((left, right) => left.time - right.time || left.order - right.order);
	const lines = ["# Content-only transcript", ""];
	if (!systemPromptCaptured) {
		lines.push(
			"Exact effective system prompts were not captured by this legacy run. Any available worker append-only prompt is included and labeled below.",
			"",
		);
	}
	for (const item of records) {
		const source = item.source ? `/${item.source}` : "";
		lines.push(`[${item.session} ${item.role}${source}]`, item.content.trimEnd(), "");
	}
	return { transcript: `${lines.join("\n").trimEnd()}\n`, recordCount: records.length, systemPromptCaptured };
}

function safeToPrune(runDirectory: string, outputFile: string): void {
	const resolved = path.resolve(runDirectory);
	const forbidden = new Set([path.parse(resolved).root, path.resolve(process.cwd()), path.resolve(os.homedir())]);
	if (forbidden.has(resolved)) throw new Error(`refusing to prune unsafe directory: ${resolved}`);
	if (path.dirname(path.resolve(outputFile)) !== resolved) {
		throw new Error("--prune requires the output file to be directly inside the run directory");
	}
	if (!fs.existsSync(path.join(resolved, "manifest.json"))) {
		throw new Error(`refusing to prune a directory without manifest.json: ${resolved}`);
	}
}

const args = process.argv.slice(2);
const prune = args.includes("--prune");
const positional = args.filter((arg) => arg !== "--prune");
if (positional.length < 1 || positional.length > 2) usage();
const runDirectory = path.resolve(positional[0]!);
const outputFile = path.resolve(positional[1] ?? path.join(runDirectory, "transcript.txt"));
if (!fs.statSync(runDirectory).isDirectory()) throw new Error(`not a directory: ${runDirectory}`);

const result = compactRun(runDirectory);
fs.mkdirSync(path.dirname(outputFile), { recursive: true });
fs.writeFileSync(outputFile, result.transcript);
if (fs.statSync(outputFile).size === 0 || result.recordCount === 0) throw new Error("generated transcript is empty");

if (prune) {
	safeToPrune(runDirectory, outputFile);
	for (const entry of fs.readdirSync(runDirectory)) {
		const target = path.join(runDirectory, entry);
		if (path.resolve(target) === outputFile) continue;
		fs.rmSync(target, { recursive: true, force: true });
	}
}

console.log(
	JSON.stringify({ output: outputFile, records: result.recordCount, systemPromptCaptured: result.systemPromptCaptured, pruned: prune }),
);