import type { ExtensionAPI, ExtensionCommandContext, 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) {
			// switchSession reloads the session and refreshes the chat window, unlike the
			// low-level setSessionFile which only swaps the file pointer. switchSession lives
			// on ExtensionCommandContext; the input handler's ctx is command-capable at runtime.
			const cmdCtx = ctx as unknown as ExtensionCommandContext;
			ctx.ui.setStatus("quick-question", undefined);
			ctx.ui.setWidget("quick-question", undefined);
			const result = await cmdCtx.switchSession(marker.returnSession, {
				withSession: async (replacementCtx) => {
					replacementCtx.ui.notify("Returned from quick session.", "info");
				},
			});
			if (result.cancelled) {
				ctx.ui.notify("Return to original session cancelled", "info");
			}
			return { action: "handled" };
		}

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