pi-ext

import { randomUUID } from "node:crypto";
import { readFile, writeFile } from "node:fs/promises";
import { join } from "node:path";
import type { SessionInfo } from "@earendil-works/pi-coding-agent";
import { AwsClient } from "aws4fetch";
import { sessionInfoFromPath } from "../archive.js";
import { type S3Config, S3Store, s3Config } from "../s3.js";
import type { HostIdentity } from "../state.js";

const decoder = new TextDecoder();

export const CATALOG_KEY = "catalog/hosts.json";
export function resolveE2eStorageConfig(
	env: NodeJS.ProcessEnv = process.env,
): S3Config | undefined {
	if (env.PI_BAK_E2E_S3 !== "1") return undefined;
	const { config, missing } = s3Config(env);
	if (!config) throw new Error(`Real S3 E2E requires ${missing.join(", ")}`);
	return config;
}

export const storageConfig = resolveE2eStorageConfig();

export function testAlias(): string {
	return `e2e-${randomUUID().slice(0, 8)}`;
}

function signedClient(config: S3Config): AwsClient {
	return new AwsClient({
		accessKeyId: config.accessKeyId,
		secretAccessKey: config.secretAccessKey,
		service: "s3",
		region: config.region,
		retries: 0,
	});
}

function objectUrl(config: S3Config, key: string): string {
	const endpoint = config.endpoint.replace(/\/+$/, "");
	const path = [config.bucket, ...key.split("/")]
		.map((part) => encodeURIComponent(part))
		.join("/");
	return `${endpoint}/${path}`;
}

export async function listKeys(
	config: S3Config,
	prefix: string,
): Promise<string[]> {
	const url = `${config.endpoint.replace(/\/+$/, "")}/${encodeURIComponent(config.bucket)}?list-type=2&max-keys=1000&prefix=${encodeURIComponent(prefix)}`;
	const response = await signedClient(config).fetch(url, { method: "GET" });
	const xml = await response.text();
	return [...xml.matchAll(/<Key>([^<]+)<\/Key>/g)].map(
		(match) => match[1] as string,
	);
}

export async function deleteKey(
	config: S3Config,
	key: string,
): Promise<number> {
	const response = await signedClient(config).fetch(objectUrl(config, key), {
		method: "DELETE",
	});
	return response.status;
}

/**
 * The bucket is shared with real backups, so tests only ever remove the
 * throwaway hosts they created themselves.
 */
export async function removeHosts(
	config: S3Config,
	hostIds: string[],
): Promise<void> {
	if (hostIds.length === 0) return;
	for (const hostId of hostIds) {
		for (const key of await listKeys(config, `sessions/${hostId}/`))
			await deleteKey(config, key);
		await deleteKey(config, `hosts/${hostId}.json`);
	}
	const store = new S3Store(config);
	for (let attempt = 0; attempt < 5; attempt++) {
		const current = await store.get(CATALOG_KEY);
		if (!current?.etag) return;
		const catalog = JSON.parse(decoder.decode(current.body)) as {
			version: number;
			hosts: HostIdentity[];
		};
		const hosts = catalog.hosts.filter(
			(host) => !hostIds.includes(host.hostId),
		);
		if (hosts.length === catalog.hosts.length) return;
		try {
			await store.put(
				CATALOG_KEY,
				`${JSON.stringify({ ...catalog, hosts })}\n`,
				{
					contentType: "application/json",
					ifMatch: current.etag,
				},
			);
			return;
		} catch {
			// Concurrent catalog writer; retry the read-modify-write cycle.
		}
	}
}

export function sessionBody(
	id: string,
	created: string,
	cwd: string,
	messages: string[],
	name?: string,
): string {
	// Entries mirror the real Pi session schema; Pi skips files it cannot parse.
	const lines = [
		JSON.stringify({
			type: "session",
			version: 3,
			id,
			timestamp: created,
			cwd,
		}),
	];
	let parentId: string | null = null;
	for (const text of messages) {
		const entryId = randomUUID().slice(0, 8);
		lines.push(
			JSON.stringify({
				type: "message",
				id: entryId,
				parentId,
				timestamp: created,
				message: { role: "user", content: [{ type: "text", text }] },
			}),
		);
		parentId = entryId;
	}
	if (name)
		lines.push(
			JSON.stringify({
				type: "session_info",
				id: randomUUID().slice(0, 8),
				parentId,
				timestamp: created,
				name,
			}),
		);
	return `${lines.join("\n")}\n`;
}

export async function writeSession(
	directory: string,
	options: {
		id?: string;
		created?: string;
		cwd?: string;
		messages?: string[];
		name?: string;
	} = {},
): Promise<SessionInfo> {
	const id = options.id ?? randomUUID();
	const created = options.created ?? new Date().toISOString();
	const path = join(directory, `${created.replace(/[:.]/g, "-")}_${id}.jsonl`);
	await writeFile(
		path,
		sessionBody(
			id,
			created,
			options.cwd ?? directory,
			options.messages ?? ["hello"],
			options.name,
		),
	);
	return sessionInfoFromPath(path);
}

export async function appendMessage(
	path: string,
	text: string,
): Promise<SessionInfo> {
	const body = await readFile(path, "utf8");
	await writeFile(
		path,
		`${body}${JSON.stringify({ type: "message", role: "assistant", content: text })}\n`,
	);
	return sessionInfoFromPath(path);
}