import { randomUUID } from "node:crypto";
import { promises as fs } from "node:fs";
import { dirname, join } from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { getAgentDir } from "@earendil-works/pi-coding-agent";
import type { AutocompleteItem } from "@earendil-works/pi-tui";
import { StringEnum } from "@earendil-works/pi-ai";
import { Type } from "typebox";

const ROOT = join(getAgentDir(), "peer-comms");
const REGISTRY_DIR = join(ROOT, "sessions");
const MAILBOX_DIR = join(ROOT, "mailboxes");
const HEARTBEAT_MS = 5_000;
const POLL_MS = 1_000;
const STALE_AFTER_MS = 15_000;
const MAX_MESSAGE_LENGTH = 10_000;

const NAME_ADJECTIVES = [
	"amber", "brisk", "calm", "clever", "cloudy", "cobalt", "cosmic", "crimson", "curious", "dapper",
	"dawn", "fuzzy", "gentle", "golden", "hidden", "jolly", "kind", "lucky", "mellow", "misty",
	"nimble", "peppy", "quiet", "rapid", "rustic", "shy", "silver", "sleepy", "solar", "steady",
	"sunny", "swift", "tidy", "velvet", "vivid", "warm", "wild", "witty", "wise", "zesty",
];
const NAME_ANIMALS = [
	"badger", "beaver", "bison", "otter", "cougar", "crane", "dingo", "eagle", "falcon", "fox",
	"gecko", "hare", "heron", "ibis", "koala", "lemur", "lynx", "moose", "newt", "otter",
	"panda", "quail", "raven", "robin", "sable", "shark", "stoat", "tiger", "viper", "wolf",
	"yak", "zebra",
];

interface LastReply {
	text: string;
	at: number;
}

interface PeerRecord {
	id: string;
	sessionId: string;
	name: string;
	cwd: string;
	pid: number;
	startedAt: number;
	heartbeatAt: number;
	status: "idle" | "working";
	lastReply?: LastReply;
}

interface PeerMessage {
	id: string;
	fromId: string;
	fromSessionId: string;
	fromName: string;
	createdAt: number;
	text: string;
}

interface Runtime {
	record: PeerRecord;
	registryPath: string;
	mailboxPath: string;
	disposed: boolean;
}

let runtime: Runtime | undefined;
let heartbeatTimer: ReturnType<typeof setInterval> | undefined;
let pollTimer: ReturnType<typeof setInterval> | undefined;
let pollPromise: Promise<void> | undefined;
let registryWrite = Promise.resolve();

function textContent(message: unknown): string {
	if (!message || typeof message !== "object") return "";
	const content = (message as { content?: unknown }).content;
	if (!Array.isArray(content)) return "";
	return content
		.filter((part): part is { type: "text"; text: string } => {
			return Boolean(part && typeof part === "object" && (part as any).type === "text" && typeof (part as any).text === "string");
		})
		.map((part) => part.text)
		.join("\n")
		.trim();
}

async function ensurePrivateDirectory(directory: string): Promise<void> {
	await fs.mkdir(directory, { recursive: true, mode: 0o700 });
	await fs.chmod(directory, 0o700);
}

async function writeJsonAtomic(filePath: string, value: unknown): Promise<void> {
	await ensurePrivateDirectory(dirname(filePath));
	const temporaryPath = `${filePath}.${process.pid}.${randomUUID()}.tmp`;
	try {
		await fs.writeFile(temporaryPath, JSON.stringify(value), { encoding: "utf8", mode: 0o600 });
		await fs.rename(temporaryPath, filePath);
	} finally {
		await fs.unlink(temporaryPath).catch(() => {});
	}
}

function persistRecord(): Promise<void> {
	const current = runtime;
	if (!current || current.disposed) return Promise.resolve();
	const snapshot = { ...current.record, lastReply: current.record.lastReply && { ...current.record.lastReply } };
	registryWrite = registryWrite
		.catch(() => {})
		.then(() => writeJsonAtomic(current.registryPath, snapshot));
	return registryWrite;
}

function stopTimer(timer: ReturnType<typeof setInterval> | undefined): void {
	if (timer) clearInterval(timer);
}

async function stopRuntime(): Promise<void> {
	stopTimer(heartbeatTimer);
	stopTimer(pollTimer);
	heartbeatTimer = undefined;
	pollTimer = undefined;

	const current = runtime;
	runtime = undefined;
	if (!current) return;
	current.disposed = true;

	await pollPromise?.catch(() => {});
	await registryWrite.catch(() => {});
	await fs.unlink(current.registryPath).catch(() => {});
	await fs.rm(current.mailboxPath, { recursive: true, force: true }).catch(() => {});
	registryWrite = Promise.resolve();
}

function nameHash(value: string): number {
	let hash = 0;
	for (const character of value) hash = (hash * 31 + character.charCodeAt(0)) >>> 0;
	return hash;
}

async function generatePeerName(sessionId: string): Promise<string> {
	const usedNames = new Set((await readPeers(true)).map((peer) => peer.name));
	const animalCount = NAME_ANIMALS.length;
	const combinations = NAME_ADJECTIVES.length * animalCount;
	const start = nameHash(sessionId) % combinations;
	for (let attempt = 0; attempt < combinations; attempt++) {
		const index = (start + attempt) % combinations;
		const name = `${NAME_ADJECTIVES[Math.floor(index / animalCount)]}-${NAME_ANIMALS[index % animalCount]}`;
		if (!usedNames.has(name)) return name;
	}
	return `${NAME_ADJECTIVES[0]}-${NAME_ANIMALS[0]}`;
}

async function startRuntime(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
	await stopRuntime();

	const sessionId = ctx.sessionManager.getSessionId() || randomUUID();
	const id = `${sessionId}-${process.pid}-${randomUUID().slice(0, 8)}`;
	const now = Date.now();
	const record: PeerRecord = {
		id,
		sessionId,
		name: await generatePeerName(sessionId),
		cwd: ctx.cwd,
		pid: process.pid,
		startedAt: now,
		heartbeatAt: now,
		status: "idle",
	};

	runtime = {
		record,
		registryPath: join(REGISTRY_DIR, `${id}.json`),
		mailboxPath: join(MAILBOX_DIR, id),
		disposed: false,
	};
	await Promise.all([
		ensurePrivateDirectory(ROOT),
		ensurePrivateDirectory(REGISTRY_DIR),
		ensurePrivateDirectory(MAILBOX_DIR),
		ensurePrivateDirectory(runtime.mailboxPath),
	]);
	await persistRecord();

	heartbeatTimer = setInterval(() => {
		const current = runtime;
		if (!current || current.disposed) return;
		current.record.heartbeatAt = Date.now();
		void persistRecord();
	}, HEARTBEAT_MS);
	heartbeatTimer.unref?.();

	pollTimer = setInterval(() => void pollMailbox(pi), POLL_MS);
	pollTimer.unref?.();
	await pollMailbox(pi);
}

function validRecord(value: unknown): value is PeerRecord {
	if (!value || typeof value !== "object") return false;
	const record = value as Partial<PeerRecord>;
	const lastReply = record.lastReply;
	return (
		typeof record.id === "string" &&
		typeof record.sessionId === "string" &&
		typeof record.name === "string" &&
		typeof record.cwd === "string" &&
		Number.isInteger(record.pid) &&
		Number.isFinite(record.startedAt) &&
		Number.isFinite(record.heartbeatAt) &&
		(record.status === "idle" || record.status === "working") &&
		(!lastReply ||
			(typeof lastReply.text === "string" && Number.isFinite(lastReply.at)))
	);
}

async function readPeers(includeStale = false): Promise<PeerRecord[]> {
	let names: string[];
	try {
		names = await fs.readdir(REGISTRY_DIR);
	} catch {
		return [];
	}

	const now = Date.now();
	const peers = await Promise.all(
		names
			.filter((name) => name.endsWith(".json"))
			.map(async (name) => {
				try {
					const value: unknown = JSON.parse(await fs.readFile(join(REGISTRY_DIR, name), "utf8"));
					return validRecord(value) ? value : undefined;
				} catch {
					return undefined;
				}
			}),
	);

	return peers
		.filter((peer): peer is PeerRecord => Boolean(peer))
		.filter((peer) => includeStale || now - peer.heartbeatAt <= STALE_AFTER_MS)
		.sort((a, b) => a.name.localeCompare(b.name) || a.id.localeCompare(b.id));
}

function isPeerMessage(value: unknown): value is PeerMessage {
	if (!value || typeof value !== "object") return false;
	const message = value as Partial<PeerMessage>;
	return (
		typeof message.id === "string" &&
		typeof message.fromId === "string" &&
		typeof message.fromSessionId === "string" &&
		typeof message.fromName === "string" &&
		Number.isFinite(message.createdAt) &&
		typeof message.text === "string" &&
		message.text.length <= MAX_MESSAGE_LENGTH
	);
}

function ageLabel(timestamp: number): string {
	const seconds = Math.max(0, Math.round((Date.now() - timestamp) / 1_000));
	return seconds === 0 ? "now" : `${seconds}s ago`;
}

function publicPeer(peer: PeerRecord): Record<string, unknown> {
	return {
		id: peer.id,
		sessionId: peer.sessionId,
		name: peer.name,
		cwd: peer.cwd,
		pid: peer.pid,
		status: peer.status,
		lastHeartbeat: ageLabel(peer.heartbeatAt),
		lastReply: peer.lastReply
			? { text: peer.lastReply.text, at: new Date(peer.lastReply.at).toISOString() }
			: null,
	};
}

function editDistance(left: string, right: string): number {
	const previous = Array.from({ length: right.length + 1 }, (_, index) => index);
	for (let i = 1; i <= left.length; i++) {
		let diagonal = previous[0];
		previous[0] = i;
		for (let j = 1; j <= right.length; j++) {
			const above = previous[j];
			previous[j] = left[i - 1] === right[j - 1]
				? diagonal
				: 1 + Math.min(diagonal, previous[j], previous[j - 1]);
			diagonal = above;
		}
	}
	return previous[right.length];
}

function resolvePeer(target: string, peers: PeerRecord[]): PeerRecord {
	const needle = target.trim().toLowerCase();
	if (!needle) throw new Error("A target session name or id is required.");

	const exact = peers.filter((peer) =>
		[peer.id, peer.sessionId, peer.name].some((value) => value.toLowerCase() === needle),
	);
	if (exact.length === 1) return exact[0];
	if (exact.length > 1) throw new Error(`Target is ambiguous: ${exact.map((peer) => peer.name).join(", ")}`);

	const prefixes = peers.filter((peer) =>
		[peer.id, peer.sessionId, peer.name].some((value) => value.toLowerCase().startsWith(needle)),
	);
	if (prefixes.length === 1) return prefixes[0];
	if (prefixes.length > 1) {
		throw new Error(`Target is ambiguous: ${prefixes.map((peer) => `${peer.name} (${peer.id})`).join(", ")}`);
	}

	const ranked = peers
		.map((peer) => ({ peer, distance: editDistance(needle, peer.name.toLowerCase()) }))
		.sort((a, b) => a.distance - b.distance);
	const closest = ranked[0];
	const next = ranked[1];
	const threshold = Math.max(2, Math.floor(needle.length * 0.35));
	if (closest && closest.distance <= threshold && (!next || next.distance > closest.distance)) return closest.peer;
	throw new Error(`No active Pi session matches "${target.trim()}". Use tab completion or action=list.`);
}

async function sendPeerMessage(target: PeerRecord, text: string): Promise<string> {
	const current = runtime;
	if (!current) throw new Error("This Pi session is not registered yet.");
	if (target.id === current.record.id) throw new Error("The target must be another Pi session.");
	if (!text.trim()) throw new Error("Message cannot be empty.");
	if (text.length > MAX_MESSAGE_LENGTH) throw new Error(`Message is too long (max ${MAX_MESSAGE_LENGTH} characters).`);

	const message: PeerMessage = {
		id: randomUUID(),
		fromId: current.record.id,
		fromSessionId: current.record.sessionId,
		fromName: current.record.name,
		createdAt: Date.now(),
		text,
	};
	const filePath = join(MAILBOX_DIR, target.id, `${message.createdAt}-${message.id}.json`);
	await writeJsonAtomic(filePath, message);
	return message.id;
}

async function pollMailbox(pi: ExtensionAPI): Promise<void> {
	const current = runtime;
	if (!current || current.disposed || pollPromise) return;
	pollPromise = (async () => {
	try {
		const names = await fs.readdir(current.mailboxPath);
		for (const name of names.filter((name) => name.includes(".json.processing-"))) {
			const claimedPath = join(current.mailboxPath, name);
			const age = Date.now() - (await fs.stat(claimedPath).then((stat) => stat.mtimeMs).catch(() => Date.now()));
			if (age >= STALE_AFTER_MS) {
				const restoredPath = claimedPath.slice(0, claimedPath.indexOf(".processing-"));
				await fs.rename(claimedPath, restoredPath).catch(() => {});
			}
		}

		const messageNames = (await fs.readdir(current.mailboxPath)).filter((name) => name.endsWith(".json")).sort();
		for (const name of messageNames) {
			const sourcePath = join(current.mailboxPath, name);
			const claimedPath = `${sourcePath}.processing-${process.pid}`;
			try {
				await fs.rename(sourcePath, claimedPath);
			} catch {
				continue;
			}

			let delivered = false;
			try {
				const value: unknown = JSON.parse(await fs.readFile(claimedPath, "utf8"));
				if (isPeerMessage(value)) {
					const message = value;
					pi.sendMessage(
						{
							customType: "peer-comms",
							content: `Peer message from ${message.fromName}:\n\n${message.text}`,
							display: true,
							details: {
								messageId: message.id,
								fromId: message.fromId,
								fromSessionId: message.fromSessionId,
								createdAt: message.createdAt,
							},
						},
						{ deliverAs: "steer", triggerTurn: true },
					);
					delivered = true;
				}
			} catch {
				// Leave a transient delivery failure in the mailbox for the next poll.
			}

			if (delivered) {
				await fs.unlink(claimedPath).catch(() => {});
			} else {
				await fs.rename(claimedPath, sourcePath).catch(() => {});
			}
		}
	} catch {
		// The mailbox can disappear during /new, /resume, or /reload.
	} finally {
		pollPromise = undefined;
	}
	})();
	await pollPromise;
}

async function setStatus(status: PeerRecord["status"]): Promise<void> {
	if (!runtime || runtime.disposed) return;
	runtime.record.status = status;
	runtime.record.heartbeatAt = Date.now();
	await persistRecord();
}

async function saveLastReply(message: unknown): Promise<void> {
	const text = textContent(message);
	if (!text || !runtime || runtime.disposed) return;
	runtime.record.lastReply = { text: text.slice(-4_000), at: Date.now() };
	runtime.record.heartbeatAt = Date.now();
	await persistRecord();
}

function formatPeers(peers: PeerRecord[]): string {
	if (peers.length === 0) return "No active Pi sessions found.";
	return peers
		.map((peer) => {
			const reply = peer.lastReply ? ` | last reply: ${peer.lastReply.text.slice(0, 120)}` : "";
			return `${peer.name} [${peer.id}] — ${peer.status}, heartbeat ${ageLabel(peer.heartbeatAt)}${reply}`;
		})
		.join("\n");
}

async function peerCompletions(prefix: string): Promise<AutocompleteItem[] | null> {
	const input = prefix.trimStart();
	if (/\s/.test(input)) return null;
	const query = input.toLowerCase();
	const peers = await readPeers();
	const items = peers
		.map((peer) => ({
			peer,
			score: peer.name.toLowerCase().startsWith(query) ? 0 : peer.name.toLowerCase().includes(query) ? 1 : 2,
		}))
		.filter(({ score }) => !query || score < 2)
		.sort((a, b) => a.score - b.score || a.peer.name.localeCompare(b.peer.name))
		.map(({ peer }) => ({
			value: peer.name,
			label: peer.name,
			description: `${peer.status} · ${peer.cwd}`,
		}));
	return items.length > 0 ? items : null;
}

function parseSendArgs(args: string): { target: string; message: string } {
	const input = args.trim();
	const quoted = input.match(/^(?:"([^"]+)"|'([^']+)')\s+([\s\S]+)$/);
	if (quoted) return { target: quoted[1] || quoted[2], message: quoted[3].trim() };
	const match = input.match(/^(\S+)\s+([\s\S]+)$/);
	if (!match) throw new Error("Usage: /peer-send <session-id-or-name> <message>");
	return { target: match[1], message: match[2].trim() };
}

const peerToolParameters = Type.Object({
	action: StringEnum(["list", "send", "read"] as const),
	target: Type.Optional(Type.String({ description: "Session id, session name, or unique id prefix" })),
	message: Type.Optional(Type.String({ description: "Message to deliver to the target session" })),
	includeStale: Type.Optional(Type.Boolean({ description: "Include sessions whose heartbeat has expired" })),
});

export default function peerCommsExtension(pi: ExtensionAPI) {
	async function ensureRuntime(ctx: ExtensionContext): Promise<void> {
		if (!runtime) await startRuntime(pi, ctx);
	}

	pi.registerTool({
		name: "peer_sessions",
		label: "Peer Sessions",
		description: "List other running Pi sessions, send them a message, or read their latest reply and status.",
		promptSnippet: "List running Pi sessions, send peer messages, or read peer status",
		promptGuidelines: [
			"Use peer_sessions action=list before addressing a peer by name when the target is unknown.",
			"Use peer_sessions action=send to ask another running Pi session to do or report something.",
			"Use peer_sessions action=read to inspect the target's latest assistant reply and current status.",
		],
		parameters: peerToolParameters,
		async execute(_toolCallId, params, _signal, _onUpdate, ctx) {
			await ensureRuntime(ctx);
			const peers = await readPeers(params.includeStale ?? false);

			if (params.action === "list") {
				return { content: [{ type: "text", text: JSON.stringify(peers.map(publicPeer), null, 2) }], details: { peers } };
			}

			const target = params.target || (runtime ? runtime.record.id : "");
			const peer = resolvePeer(target, peers);
			if (params.action === "send") {
				if (!params.message) throw new Error("message is required for action=send");
				const messageId = await sendPeerMessage(peer, params.message);
				return {
					content: [{ type: "text", text: `Message ${messageId} queued for ${peer.name} (${peer.id}).` }],
					details: { peer: publicPeer(peer), messageId },
				};
			}

			const result = {
				...publicPeer(peer),
				latestReply: peer.lastReply?.text || null,
			};
			return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }], details: { peer } };
		},
	});

	pi.registerCommand("peers", {
		description: "List active Pi peer sessions",
		handler: async (_args, ctx) => {
			await ensureRuntime(ctx);
			ctx.ui.notify(formatPeers(await readPeers()), "info");
		},
	});

	pi.registerCommand("peer-send", {
		description: "Send a message to another Pi session",
		getArgumentCompletions: peerCompletions,
		handler: async (args, ctx) => {
			try {
				await ensureRuntime(ctx);
				const { target, message } = parseSendArgs(args);
				const peers = await readPeers();
				const peer = resolvePeer(target, peers);
				const messageId = await sendPeerMessage(peer, message);
				ctx.ui.notify(`Queued ${messageId} for ${peer.name}.`, "info");
			} catch (error) {
				ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
			}
		},
	});

	pi.registerCommand("peer-read", {
		description: "Read a peer's latest reply and status",
		getArgumentCompletions: peerCompletions,
		handler: async (args, ctx) => {
			try {
				await ensureRuntime(ctx);
				const peers = await readPeers();
				const peer = resolvePeer(args.trim() || runtime?.record.id || "", peers);
				const reply = peer.lastReply?.text || "(no assistant reply recorded yet)";
				ctx.ui.notify(`${peer.name} — ${peer.status}\n\n${reply}`, "info");
			} catch (error) {
				ctx.ui.notify(error instanceof Error ? error.message : String(error), "error");
			}
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		await startRuntime(pi, ctx);
	});

	pi.on("agent_start", async () => {
		await setStatus("working");
	});

	pi.on("agent_settled", async () => {
		await setStatus("idle");
	});

	pi.on("message_end", async (event) => {
		if (event.message.role === "assistant") await saveLastReply(event.message);
	});

	pi.on("session_shutdown", async () => {
		await stopRuntime();
	});
}