import { ExtensionRunner } from "@earendil-works/pi-coding-agent";
import type { ExtensionCommandContext, ExtensionContext } from "@earendil-works/pi-coding-agent";

let navigateTreeFn: ExtensionCommandContext["navigateTree"] | null = null;
let waitForIdleFn: ExtensionCommandContext["waitForIdle"] | null = null;
let captureInstalled = false;

export function installCommandActionCapture(): void {
	if (captureInstalled) {
		return;
	}
	captureInstalled = true;

	const original = ExtensionRunner.prototype.bindCommandContext;
	ExtensionRunner.prototype.bindCommandContext = function bindCommandContext(actions) {
		if (actions?.navigateTree && actions?.waitForIdle) {
			navigateTreeFn = actions.navigateTree;
			waitForIdleFn = actions.waitForIdle;
		} else {
			navigateTreeFn = null;
			waitForIdleFn = null;
		}
		return original.call(this, actions);
	};
}

export function cacheCommandActions(ctx: ExtensionCommandContext): void {
	navigateTreeFn = ctx.navigateTree.bind(ctx);
	waitForIdleFn = ctx.waitForIdle.bind(ctx);
}

export function hasCachedCommandActions(): boolean {
	return navigateTreeFn !== null && waitForIdleFn !== null;
}

export function makeCommandContext(ctx: ExtensionContext): ExtensionCommandContext | null {
	if (!navigateTreeFn || !waitForIdleFn) {
		return null;
	}

	return {
		...ctx,
		waitForIdle: waitForIdleFn,
		navigateTree: navigateTreeFn,
		newSession: async () => ({ cancelled: true }),
		fork: async () => ({ cancelled: true }),
		switchSession: async () => ({ cancelled: true }),
		reload: async () => {},
		getSystemPromptOptions: () => ({ cwd: ctx.cwd }),
	} as ExtensionCommandContext;
}

export function __resetCommandActionsForTests(): void {
	navigateTreeFn = null;
	waitForIdleFn = null;
}