pi-ext

import type {
	ExtensionContext,
	SessionInfo,
} from "@earendil-works/pi-coding-agent";
import { SessionManager } from "@earendil-works/pi-coding-agent";
import type { BakArchive } from "./archive.js";
import { type S3ConfigResolution, s3Config } from "./config.js";
import { closeDebug, dbg } from "./src/debug.ts";
import {
	type CatalogCache,
	clearPending,
	type HostIdentity,
	markPending,
	type PendingEntry,
	readPending,
	type SessionRecord,
	statePaths,
} from "./state.js";

const ICON = "󰁯";
const STATUS_KEY = "bak";
const SHUTDOWN_DRAIN_MS = 2_000;
const SUBCOMMANDS = [
	"init",
	"backup",
	"restore",
	"search",
	"refresh",
	"status",
] as const;
const SEARCH_LIMIT = 10;

type ArchiveModule = typeof import("./archive.js");

interface Backend {
	archive: BakArchive;
	archiveModule: ArchiveModule;
}

function currentSession(ctx: ExtensionContext): SessionInfo | undefined {
	const path = ctx.sessionManager.getSessionFile();
	const header = ctx.sessionManager.getHeader();
	if (!path || !header) return undefined;
	return {
		path,
		id: header.id,
		cwd: header.cwd,
		name: ctx.sessionManager.getSessionName(),
		parentSessionPath: header.parentSession,
		created: new Date(header.timestamp),
		modified: new Date(),
		messageCount: ctx.sessionManager
			.getEntries()
			.filter((entry) => entry.type === "message").length,
		firstMessage: "",
		allMessagesText: "",
	};
}

function errorText(error: unknown): string {
	const text = error instanceof Error ? error.message : String(error);
	return text.replace(/\s+/g, " ").trim().slice(0, 240);
}

function formatAge(timestamp: string | undefined): string {
	if (!timestamp) return "never";
	const age = Date.now() - new Date(timestamp).getTime();
	if (!Number.isFinite(age) || age < 0) return "unknown";
	if (age < 60_000) return "now";
	if (age < 3_600_000) return `${Math.floor(age / 60_000)}m ago`;
	if (age < 86_400_000) return `${Math.floor(age / 3_600_000)}h ago`;
	return `${Math.floor(age / 86_400_000)}d ago`;
}

function destinationArgs(args: string): { target: string; path?: string } {
	const match = /^(\S+)(?:\s+(.+))?$/.exec(args.trim());
	return { target: match?.[1] ?? "", path: match?.[2]?.trim() || undefined };
}

function selectRecords(cache: CatalogCache, target: string): SessionRecord[] {
	if (target === "all") return cache.indexes.flatMap((index) => index.sessions);
	const slash = target.indexOf("/");
	const alias = slash >= 0 ? target.slice(0, slash) : undefined;
	const prefix = slash >= 0 ? target.slice(slash + 1) : target;
	if (!prefix) return [];
	return cache.indexes
		.filter((index) => !alias || index.alias === alias)
		.flatMap((index) => index.sessions)
		.filter((record) => record.id.startsWith(prefix));
}

export function searchRecords(
	cache: CatalogCache,
	query: string,
): { alias: string; record: SessionRecord; prefix: string }[] {
	const needle = query.trim().toLowerCase();
	if (!needle) return [];
	const matches: { alias: string; record: SessionRecord; ids: string[] }[] = [];
	for (const index of cache.indexes) {
		const ids = index.sessions.map((record) => record.id);
		for (const record of index.sessions) {
			if (
				!record.name?.toLowerCase().includes(needle) &&
				!record.cwd.toLowerCase().includes(needle)
			)
				continue;
			matches.push({ alias: index.alias, record, ids });
		}
	}
	matches.sort((a, b) => b.record.created.localeCompare(a.record.created));
	return matches.slice(0, SEARCH_LIMIT).map(({ alias, record, ids }) => ({
		alias,
		record,
		prefix: uniquePrefix(record.id, ids),
	}));
}

function uniquePrefix(id: string, ids: string[]): string {
	for (let length = Math.min(8, id.length); length < id.length; length++) {
		const prefix = id.slice(0, length);
		if (
			!ids.some((candidate) => candidate !== id && candidate.startsWith(prefix))
		)
			return prefix;
	}
	return id;
}

async function localSessions(ctx: ExtensionContext): Promise<SessionInfo[]> {
	const manager = ctx.sessionManager as typeof ctx.sessionManager & {
		usesDefaultSessionDir?: () => boolean;
	};
	// Non-persistent sessions report an empty directory; scanning it would read
	// the working directory instead of the session store.
	const directory = ctx.sessionManager.getSessionDir();
	return !directory || manager.usesDefaultSessionDir?.()
		? SessionManager.listAll()
		: SessionManager.listAll(directory);
}

export function createBakRuntime(
	resolved: S3ConfigResolution = s3Config(),
	missingWarningShown = false,
) {
	const paths = statePaths();
	let backendPromise: Promise<Backend> | undefined;
	const loadBackend = (): Promise<Backend | undefined> => {
		const config = resolved.config;
		if (!config) return Promise.resolve(undefined);
		if (backendPromise) return backendPromise;
		backendPromise = Promise.all([
			import("./archive.js"),
			import("./s3.js"),
		]).then(([archiveModule, { S3Store }]) => ({
			archive: new archiveModule.BakArchive(new S3Store(config), paths),
			archiveModule,
		}));
		return backendPromise;
	};
	let identity: HostIdentity | undefined;
	let cache: CatalogCache | undefined;
	let lastBackupAt: string | undefined;
	let lastError: string | undefined;
	let warnedFailure = false;
	let warnedSetup = missingWarningShown;
	let lifecycle = 0;
	let loadedLifecycle = 0;
	let backgroundLifecycle = 0;
	let serial: Promise<void> = Promise.resolve();
	let readiness: Promise<void> = Promise.resolve();
	let lifecycleController: AbortController | undefined;
	let pending = new Map<string, PendingEntry>();

	const notify = (
		ctx: ExtensionContext,
		message: string,
		level: "info" | "warning" | "error" = "info",
	) => {
		if (ctx.hasUI) ctx.ui.notify(`bak: ${message}`, level);
	};

	const runSerial = <T>(operation: () => Promise<T>): Promise<T> => {
		const result = serial.then(operation, operation);
		serial = result.then(
			() => undefined,
			() => undefined,
		);
		return result;
	};

	const requireArchive = async (
		ctx: ExtensionContext,
	): Promise<BakArchive | undefined> => {
		const backend = await loadBackend();
		if (backend) return backend.archive;
		notify(ctx, `disabled; missing ${resolved.missing.join(", ")}`, "error");
		return undefined;
	};

	const setProgress = (ctx: ExtensionContext, text?: string) => {
		if (ctx.hasUI)
			ctx.ui.setStatus(
				STATUS_KEY,
				text ? ctx.ui.theme.fg("dim", `${ICON} ${text}`) : undefined,
			);
	};

	const addPending = async (path: string): Promise<PendingEntry> => {
		const entry = await markPending(paths.pendingDir, path);
		pending.set(entry.id, entry);
		return entry;
	};

	const removePending = async (entries: PendingEntry[]) => {
		await Promise.all(
			entries.map(async (entry) => {
				await clearPending(paths.pendingDir, entry.id);
				pending.delete(entry.id);
			}),
		);
	};

	const ensureLoaded = (
		ctx: ExtensionContext,
		epoch = lifecycle,
	): Promise<void> => {
		if (!resolved.config || loadedLifecycle === epoch) return readiness;
		loadedLifecycle = epoch;
		readiness = (async () => {
			const backend = await loadBackend();
			if (!backend) return;
			const [savedIdentity, savedPending] = await Promise.all([
				backend.archive.identity(),
				readPending(paths.pendingDir),
			]);
			let savedCache: CatalogCache | undefined;
			try {
				savedCache = await backend.archive.cachedCatalog();
			} catch {
				savedCache = undefined;
			}
			if (epoch !== lifecycle) return;
			identity = savedIdentity;
			cache = savedCache;
			pending = new Map(savedPending.map((entry) => [entry.id, entry]));
			if (!identity && !warnedSetup) {
				notify(ctx, "run /bak init <alias>", "warning");
				warnedSetup = true;
			}
		})();
		return readiness;
	};

	const refresh = async (
		ctx: ExtensionContext,
		signal?: AbortSignal,
	): Promise<CatalogCache | undefined> => {
		const client = await requireArchive(ctx);
		if (!client) return undefined;
		const next = await client.refresh(signal);
		cache = next;
		return next;
	};

	const backupBatch = async (
		infos: SessionInfo[],
		onProgress?: (done: number, total: number) => void,
		replayEntries?: PendingEntry[],
	) => {
		if (infos.length === 0) return { uploaded: 0, indexed: 0 };
		const backend = await loadBackend();
		if (!backend) return { uploaded: 0, indexed: 0 };
		const archive = backend.archive;
		await readiness;
		if (!replayEntries)
			await Promise.all(infos.map((info) => addPending(info.path)));
		const paths = new Set(infos.map((info) => info.path));
		const covered = (replayEntries ?? [...pending.values()]).filter((entry) =>
			paths.has(entry.path),
		);
		if (!identity) throw new Error("Run /bak init <alias> first");
		try {
			const result = await archive.backup(
				infos,
				onProgress,
				lifecycleController?.signal,
			);
			await removePending(covered);
			lastBackupAt = new Date().toISOString();
			lastError = undefined;
			return result;
		} catch (error) {
			if (error instanceof backend.archiveModule.PartialBackupError) {
				const succeeded = new Set(error.succeededPaths);
				await removePending(
					covered.filter((entry) => succeeded.has(entry.path)),
				);
			}
			throw error;
		}
	};

	const startBackground = (ctx: ExtensionContext) => {
		const epoch = lifecycle;
		if (!resolved.config || backgroundLifecycle === epoch) return;
		backgroundLifecycle = epoch;
		void ensureLoaded(ctx, epoch)
			.then(async () => {
				if (!identity || epoch !== lifecycle) return;
				const backend = await loadBackend();
				if (!backend) return;
				try {
					await backend.archive.ensureInitialized(
						identity,
						lifecycleController?.signal,
					);
				} catch (error) {
					if (epoch === lifecycle)
						notify(
							ctx,
							`background init failed: ${errorText(error)}`,
							"warning",
						);
				}
				void refresh(ctx, lifecycleController?.signal).catch((error) => {
					if (epoch === lifecycle && !cache)
						notify(
							ctx,
							`catalog refresh failed: ${errorText(error)}`,
							"warning",
						);
				});
				await retryPending(epoch);
			})
			.catch((error) => {
				if (epoch === lifecycle)
					notify(
						ctx,
						`background state failed: ${errorText(error)}`,
						"warning",
					);
			});
	};

	const automaticBackup = (ctx: ExtensionContext) => {
		startBackground(ctx);
		const info = currentSession(ctx);
		if (!resolved.config || !info) {
			dbg?.("operation.finish", {
				operation: "automatic-backup",
				outcome: "skipped",
			});
			return;
		}
		const epoch = lifecycle;
		void runSerial(async () => {
			const statusTimer = setTimeout(() => {
				if (epoch === lifecycle) setProgress(ctx, "↑");
			}, 300);
			try {
				await backupBatch([info]);
				dbg?.("operation.finish", {
					operation: "automatic-backup",
					outcome: "success",
				});
				if (warnedFailure && epoch === lifecycle) {
					notify(ctx, "backup recovered");
					warnedFailure = false;
				}
			} catch (error) {
				dbg?.("operation.finish", {
					operation: "automatic-backup",
					outcome: "failed",
				});
				lastError = errorText(error);
				if (!warnedFailure && epoch === lifecycle) {
					notify(ctx, `backup failed: ${lastError}; retry pending`, "warning");
					warnedFailure = true;
				}
			} finally {
				clearTimeout(statusTimer);
				if (epoch === lifecycle) setProgress(ctx);
			}
		});
	};

	const retryPending = async (epoch: number) => {
		if (!resolved.config || pending.size === 0 || epoch !== lifecycle) return;
		const backend = await loadBackend();
		if (!backend) return;
		const replayEntries = [...pending.values()];
		const infos: SessionInfo[] = [];
		for (const path of new Set(replayEntries.map((entry) => entry.path))) {
			try {
				infos.push(await backend.archiveModule.sessionInfoFromPath(path));
			} catch (error) {
				lastError = `pending ${path}: ${errorText(error)}`;
			}
		}
		if (infos.length === 0) return;
		await runSerial(() => backupBatch(infos, undefined, replayEntries));
	};

	const sessionStart = (ctx: ExtensionContext) => {
		lifecycle++;
		lifecycleController?.abort();
		lifecycleController = new AbortController();
		loadedLifecycle = 0;
		backgroundLifecycle = 0;
		readiness = Promise.resolve();
		identity = undefined;
		cache = undefined;
		pending = new Map();
		if (!resolved.config && !warnedSetup) {
			notify(
				ctx,
				`disabled; missing ${resolved.missing.join(", ")}`,
				"warning",
			);
			warnedSetup = true;
		}
	};

	/** With automatic=false no upload starts; an explicit /bak backup still drains. */
	const sessionShutdown = async (ctx: ExtensionContext, automatic = true) => {
		const info = currentSession(ctx);
		if (automatic && resolved.config && info) {
			await addPending(info.path).catch((error) => {
				lastError = errorText(error);
			});
		}
		if (automatic) automaticBackup(ctx);
		await Promise.race([
			serial,
			new Promise((resolve) => setTimeout(resolve, SHUTDOWN_DRAIN_MS)),
		]);
		lifecycleController?.abort();
		lifecycle++;
		setProgress(ctx);
	};

	const command = {
		getArgumentCompletions(prefix: string) {
			const normalized = prefix.trimStart();
			const restore = /^restore\s+(\S*)$/i.exec(normalized);
			if (restore) {
				const query = restore[1]?.toLowerCase() ?? "";
				const records =
					cache?.indexes.flatMap((index) => {
						const ids = index.sessions.map((record) => record.id);
						return index.sessions.map((record) => {
							const prefix = uniquePrefix(record.id, ids);
							return {
								value: `restore ${index.alias}/${prefix}`,
								label: `${index.alias}/${prefix}`,
								description: record.name ?? record.cwd,
							};
						});
					}) ?? [];
				const items = [
					{ value: "restore all", label: "all", description: "all sessions" },
					...records,
				].filter((item) => item.label.toLowerCase().startsWith(query));
				return items.length > 0 ? items.slice(0, 100) : null;
			}
			if (normalized.includes(" ")) return null;
			const items = SUBCOMMANDS.filter((command) =>
				command.startsWith(normalized.toLowerCase()),
			).map((command) => ({ value: command, label: command }));
			return items.length > 0 ? items : null;
		},
		async handler(args: string, ctx: ExtensionContext) {
			const trimmed = args.trim();
			const [command = "status", ...rest] = trimmed.split(/\s+/);
			const subcommand =
				SUBCOMMANDS.find((candidate) => candidate === command) ?? "invalid";
			let outcome: "success" | "failed" | "skipped" = "success";
			try {
				if (command === "status") {
					if (!resolved.config) outcome = "skipped";
					else
						await ensureLoaded(ctx).catch((error) => {
							outcome = "failed";
							lastError = errorText(error);
						});
					const indexed = cache?.indexes.reduce(
						(total, index) => total + index.sessions.length,
						0,
					);
					const state = resolved.config
						? identity
							? "ready"
							: "not initialized"
						: `disabled; missing ${resolved.missing.join(", ")}`;
					notify(
						ctx,
						[
							identity?.alias,
							state,
							indexed === undefined ? undefined : `${indexed} indexed`,
							`${new Set([...pending.values()].map((entry) => entry.path)).size} pending`,
							`last backup ${formatAge(lastBackupAt)}`,
							`catalog ${formatAge(cache?.refreshedAt)}`,
							lastError ? `error ${lastError}` : undefined,
						]
							.filter(Boolean)
							.join(" · "),
						lastError ? "warning" : "info",
					);
					return;
				}
				const client = await requireArchive(ctx);
				if (!client) {
					outcome = "skipped";
					return;
				}
				await ensureLoaded(ctx);
				if (command === "init") {
					if (rest.length !== 1) throw new Error("Usage: /bak init <alias>");
					setProgress(ctx, "initializing");
					identity = await runSerial(() =>
						client.init(rest[0] ?? "", lifecycleController?.signal),
					);
					cache = await refresh(ctx, lifecycleController?.signal);
					warnedSetup = false;
					void retryPending(lifecycle).catch((error) => {
						lastError = errorText(error);
						notify(ctx, `pending retry failed: ${lastError}`, "warning");
					});
					notify(ctx, `initialized as ${identity.alias}`);
					return;
				}
				if (command === "refresh") {
					if (rest.length > 0) throw new Error("Usage: /bak refresh");
					setProgress(ctx, "refreshing");
					const next = await runSerial(() =>
						client.refresh(lifecycleController?.signal),
					);
					cache = next;
					notify(ctx, `refreshed ${next.indexes.length} host(s)`);
					return;
				}
				if (command === "backup") {
					if (rest.length > 0) throw new Error("Usage: /bak backup");
					setProgress(ctx, "scanning");
					const infos = await localSessions(ctx);
					const result = await runSerial(() =>
						backupBatch(infos, (done, total) =>
							setProgress(ctx, `↑ ${done}/${total}`),
						),
					);
					notify(
						ctx,
						`backed up ${result.uploaded}; ${result.indexed} indexed`,
					);
					return;
				}
				if (command === "search") {
					const query = rest.join(" ").trim();
					if (!query) throw new Error("Usage: /bak search <name or path>");
					if (!cache) {
						setProgress(ctx, "refreshing");
						cache = await runSerial(() =>
							client.refresh(lifecycleController?.signal),
						);
					}
					const results = searchRecords(cache, query);
					notify(
						ctx,
						results.length
							? results
									.map(
										({ alias, record, prefix }) =>
											`${alias}/${prefix} ${record.name ?? record.cwd} (${record.created.slice(0, 10)})`,
									)
									.join("\n") +
									(results.length === SEARCH_LIMIT
										? "\nShowing first 10; narrow query for more"
										: "")
							: `no sessions match ${query}`,
					);
					return;
				}
				if (command === "restore") {
					const parsed = destinationArgs(rest.join(" "));
					if (!parsed.target)
						throw new Error(
							"Usage: /bak restore <alias>/<id-prefix>|all [path]",
						);
					setProgress(ctx, "refreshing");
					const current = await runSerial(() =>
						client.refresh(lifecycleController?.signal),
					);
					cache = current;
					const records = selectRecords(current, parsed.target);
					if (records.length === 0)
						throw new Error(`No session matches ${parsed.target}`);
					if (parsed.target !== "all" && records.length !== 1)
						throw new Error(`Ambiguous session: ${parsed.target}`);
					const result = await runSerial(() =>
						client.restore(
							records,
							parsed.path ?? ctx.sessionManager.getSessionDir(),
							(done, total) => setProgress(ctx, `↓ ${done}/${total}`),
							lifecycleController?.signal,
						),
					);
					notify(
						ctx,
						`restored ${result.restored}, skipped ${result.skipped}, conflicts ${result.conflicts}`,
						result.conflicts > 0 ? "warning" : "info",
					);
					return;
				}
				throw new Error(
					"Usage: /bak init <alias> | backup | restore <alias>/<id-prefix>|all [path] | search <name or path> | refresh | status",
				);
			} catch (error) {
				outcome = "failed";
				lastError = errorText(error);
				notify(ctx, lastError, "error");
			} finally {
				dbg?.("operation.finish", {
					operation: "command",
					outcome,
					subcommand,
				});
				if (command !== "status") setProgress(ctx);
			}
		},
	};

	return {
		automaticBackup,
		closeDebugWhenIdle() {
			void serial.finally(closeDebug);
		},
		command,
		sessionShutdown,
		sessionStart,
	};
}