pi-ext

import { mkdir, mkdtemp, readdir, readFile, rm } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import type { SessionInfo } from "@earendil-works/pi-coding-agent";
import { afterAll, beforeAll, describe, expect, it } from "vitest";
import { BakArchive } from "../archive.js";
import bakExtension from "../index.js";
import { type S3Config, S3Store } from "../s3.js";
import { type HostIdentity, statePaths } from "../state.js";
import {
	removeHosts,
	storageConfig,
	testAlias,
	writeSession,
} from "./support.js";

const config = storageConfig;

interface Note {
	message: string;
	level: string;
}

type CommandHandler = (args: string, ctx: unknown) => Promise<void> | void;
type EventHandler = (event: unknown, ctx: unknown) => Promise<void> | void;

const aliases: string[] = [];
let workspace: string;
let sessionDir: string;
let notes: Note[] = [];
let run: (args: string) => Promise<void>;
let emit: (event: string) => Promise<void>;
let context: Record<string, unknown>;
let currentSession: SessionInfo | undefined;

/**
 * Registers the extension exactly like Pi does, but against a throwaway state
 * directory and session store so real backups stay untouched.
 */
function loadExtension(): void {
	const handlers = new Map<string, EventHandler[]>();
	let command: { handler: CommandHandler } | undefined;
	const pi = {
		on(event: string, handler: EventHandler) {
			handlers.set(event, [...(handlers.get(event) ?? []), handler]);
		},
		registerCommand(_name: string, definition: { handler: CommandHandler }) {
			command = definition;
		},
	};
	bakExtension(pi as never);
	if (!command) throw new Error("bak command was not registered");
	const registered = command;
	run = async (args: string) => {
		await registered.handler(args, context);
	};
	emit = async (event: string) => {
		for (const handler of handlers.get(event) ?? [])
			await handler({ type: event }, context);
	};
}

function sessionContext(): Record<string, unknown> {
	return {
		hasUI: true,
		ui: {
			notify: (message: string, level = "info") => {
				notes.push({ message, level });
			},
			setStatus: () => undefined,
			theme: { fg: (_style: string, text: string) => text },
		},
		sessionManager: {
			getSessionDir: () => sessionDir,
			usesDefaultSessionDir: () => false,
			getSessionFile: () => currentSession?.path,
			getHeader: () =>
				currentSession
					? {
							id: currentSession.id,
							cwd: currentSession.cwd,
							timestamp: currentSession.created.toISOString(),
						}
					: undefined,
			getSessionName: () => currentSession?.name,
			getEntries: () => [],
		},
	};
}

function lastNote(): Note {
	const note = notes.at(-1);
	if (!note) throw new Error("no bak notification recorded");
	return note;
}

async function archivedIds(alias: string): Promise<string[]> {
	const cache = await sandboxArchive().refresh();
	return (
		cache.indexes
			.find((index) => index.alias === alias)
			?.sessions.map((record) => record.id) ?? []
	);
}

async function waitForArchived(alias: string, id: string): Promise<void> {
	for (let attempt = 0; attempt < 30; attempt++) {
		if ((await archivedIds(alias)).includes(id)) return;
		await new Promise((resolve) => setTimeout(resolve, 500));
	}
	throw new Error(`session ${id} never reached the archive`);
}

function sandboxArchive(): BakArchive {
	return new BakArchive(
		new S3Store(config as S3Config),
		statePaths(process.env),
	);
}

describe.skipIf(!config)("bak commands against real S3 storage", () => {
	const original = {
		state: process.env.XDG_STATE_HOME,
		cache: process.env.XDG_CACHE_HOME,
	};

	beforeAll(async () => {
		workspace = await mkdtemp(join(tmpdir(), "bak-cmd-e2e-"));
		process.env.XDG_STATE_HOME = join(workspace, "state");
		process.env.XDG_CACHE_HOME = join(workspace, "cache");
		sessionDir = join(workspace, "sessions");
		await mkdir(sessionDir, { recursive: true });
		context = sessionContext();
		loadExtension();
	});

	afterAll(async () => {
		if (config && aliases.length > 0) {
			const catalog = await new S3Store(config).get("catalog/hosts.json");
			const hosts = catalog
				? (
						JSON.parse(new TextDecoder().decode(catalog.body)) as {
							hosts: HostIdentity[];
						}
					).hosts
				: [];
			await removeHosts(
				config,
				hosts
					.filter((host) => aliases.includes(host.alias))
					.map((host) => host.hostId),
			);
		}
		process.env.XDG_STATE_HOME = original.state;
		process.env.XDG_CACHE_HOME = original.cache;
		if (original.state === undefined) delete process.env.XDG_STATE_HOME;
		if (original.cache === undefined) delete process.env.XDG_CACHE_HOME;
		await rm(workspace, { recursive: true, force: true });
	});

	it("reports an uninitialized host", async () => {
		notes = [];
		await run("status");

		expect(lastNote().message).toContain("not initialized");
		expect(lastNote().message).toContain("0 pending");
	});

	it("initializes, backs up local sessions, and reports status", async () => {
		const alias = testAlias();
		aliases.push(alias);
		notes = [];

		await run(`init ${alias}`);
		expect(lastNote().message).toBe(`bak: initialized as ${alias}`);
		const identity = JSON.parse(
			await readFile(statePaths(process.env).identity, "utf8"),
		) as HostIdentity;
		expect(identity.alias).toBe(alias);

		const first = await writeSession(sessionDir, { name: "command session" });
		const second = await writeSession(sessionDir, { messages: ["x", "y"] });
		await run("backup");
		expect(lastNote().message).toBe("bak: backed up 2; 2 indexed");

		await run("status");
		expect(lastNote().message).toContain(alias);
		expect(lastNote().message).toContain("ready");
		// The counter spans every catalog host, including unrelated real hosts.
		expect(lastNote().message).toMatch(/[1-9]\d* indexed/);
		expect(lastNote().message).toContain("0 pending");
		expect(lastNote().message).not.toContain("error");

		const cache = await sandboxArchive().refresh();
		const ids =
			cache.indexes
				.find((index) => index.alias === alias)
				?.sessions.map((record) => record.id) ?? [];
		expect(ids.sort()).toEqual([first.id, second.id].sort());
	});

	it("searches remote names from the cached catalog", async () => {
		const alias = aliases.at(-1) as string;
		notes = [];
		await run("refresh");
		await run("search COMMAND SESSION");
		expect(lastNote().message).toContain(`${alias}/`);
		expect(lastNote().message).toContain("command session");
	});

	it("refreshes the catalog and restores by alias prefix", async () => {
		const alias = aliases.at(-1) as string;
		notes = [];
		await run("refresh");
		expect(lastNote().message).toMatch(/refreshed \d+ host\(s\)/);

		const cache = await sandboxArchive().refresh();
		const record = cache.indexes.find((index) => index.alias === alias)
			?.sessions[0];
		expect(record).toBeDefined();
		const target = join(workspace, "restore-target");

		await run(`restore ${alias}/${record?.id.slice(0, 8)} ${target}`);
		expect(lastNote().message).toBe("bak: restored 1, skipped 0, conflicts 0");
		expect(await readdir(target)).toHaveLength(1);

		await run(`restore ${alias}/${record?.id.slice(0, 8)} ${target}`);
		expect(lastNote().message).toBe("bak: restored 0, skipped 1, conflicts 0");
	});

	it("restores every archived session of the host", async () => {
		const alias = aliases.at(-1) as string;
		const target = join(workspace, "restore-all");
		notes = [];

		await run(`restore all ${target}`);
		expect(lastNote().message).toMatch(/restored \d+, skipped \d+/);

		const cache = await sandboxArchive().refresh();
		const own =
			cache.indexes.find((index) => index.alias === alias)?.sessions.length ??
			0;
		expect((await readdir(target)).length).toBeGreaterThanOrEqual(own);
	});

	it("keeps working after session lifecycle events", async () => {
		notes = [];
		await emit("session_start");
		await writeSession(sessionDir, { messages: ["after restart"] });
		await run("backup");
		expect(lastNote().message).toBe("bak: backed up 3; 3 indexed");

		await run("status");
		expect(lastNote().message).toMatch(/[1-9]\d* indexed/);
		expect(lastNote().message).toContain("0 pending");
		expect(lastNote().message).not.toContain("error");
	});

	it("backs up the active session when the agent settles", async () => {
		const alias = aliases.at(-1) as string;
		notes = [];
		currentSession = await writeSession(sessionDir, {
			messages: ["settled"],
			name: "settled session",
		});

		await emit("agent_settled");
		await waitForArchived(alias, currentSession.id);

		expect(notes.filter((note) => note.level !== "info")).toEqual([]);
		const markers = await readdir(statePaths(process.env).pendingDir);
		expect(markers).toEqual([]);
	});

	it("retains a pending retry after a failed backup and recovers", async () => {
		const alias = aliases.at(-1) as string;
		const bucket = process.env.PI_BAK_S3_BUCKET;
		notes = [];
		currentSession = await writeSession(sessionDir, { messages: ["pending"] });

		process.env.PI_BAK_S3_BUCKET = `${(config as S3Config).bucket}-missing`;
		loadExtension();
		await run("backup");
		expect(lastNote().level).toBe("error");
		expect(await readdir(statePaths(process.env).pendingDir)).not.toEqual([]);

		process.env.PI_BAK_S3_BUCKET = bucket;
		if (bucket === undefined) delete process.env.PI_BAK_S3_BUCKET;
		loadExtension();
		await emit("session_start");
		await emit("agent_settled");
		await waitForArchived(alias, currentSession.id);

		for (let attempt = 0; attempt < 30; attempt++) {
			if ((await readdir(statePaths(process.env).pendingDir)).length === 0)
				break;
			await new Promise((resolve) => setTimeout(resolve, 500));
		}
		expect(await readdir(statePaths(process.env).pendingDir)).toEqual([]);
	});

	it("reports usage and initialization errors", async () => {
		notes = [];
		await run("bogus");
		expect(lastNote().level).toBe("error");
		expect(lastNote().message).toContain("Usage:");

		await run(`init ${testAlias()}`);
		expect(lastNote().level).toBe("error");
		expect(lastNote().message).toMatch(/Already initialized as/);
	});
});