import type { ExtensionAPI, SessionEntry } from "@earendil-works/pi-coding-agent";

const QUICK_MARKER_TYPE = "quick-question-session";

type QuickSessionMarker = {
	returnSession: string;
};

function getQuickSessionMarker(entries: SessionEntry[]): QuickSessionMarker | undefined {
	for (let i = entries.length - 1; i >= 0; i--) {
		const entry = entries[i];
		if (entry.type !== "custom" || entry.customType !== QUICK_MARKER_TYPE) continue;

		const data = entry.data as Partial<QuickSessionMarker> | undefined;
		if (typeof data?.returnSession === "string") {
			return { returnSession: data.returnSession };
		}
	}

	return undefined;
}

export default function (pi: ExtensionAPI) {
	pi.on("input", async (event, ctx) => {
		if (event.source === "extension") return { action: "continue" };
		if (event.text.trim() !== "exit") return { action: "continue" };

		const marker = getQuickSessionMarker(ctx.sessionManager.getBranch());
		if (marker) {
			(ctx.sessionManager as unknown as { setSessionFile(path: string): void }).setSessionFile(marker.returnSession);
			ctx.ui.setStatus("quick-question", undefined);
			ctx.ui.setWidget("quick-question", undefined);
			ctx.ui.notify("Returned from quick session.", "info");
			return { action: "handled" };
		}

		ctx.shutdown();
		return { action: "handled" };
	});
}