import { getAgentDir, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";
import { readEffectiveCodexConversionConfig } from "../adapter/activation/config-store.ts";
import { syncAdapter } from "../adapter/activation/activation.ts";
import { resolveCodexRuntimePlanForState } from "../adapter/activation/runtime-plan.ts";
import { supportsCodexDeveloperMessages } from "../adapter/provider-request.ts";
import { onCodeModeExtensionToolsRefresh } from "../code-mode-extension-tools.ts";
import { extractPiPromptSkills } from "../prompt/build-system-prompt.ts";
import type { CodeModeProxyProviderRegistration } from "../providers/code-mode-proxy-provider.ts";
import { maybeWarnLocalCheckoutVersion } from "../adapter/local-version-warning.ts";
import { clearApplyPatchRenderState } from "../tools/apply-patch/tool.ts";
import type { CodeModeRegistration } from "../tools/code-mode/tools.ts";
import { initializeBashParser } from "../shell/bash.ts";
import { appendNotebookTreeEpoch, notebookSessionIdentity } from "../tools/notebook-mode/session-identity.ts";
import { activateNotebookCheckpoint, type NotebookCheckpointActivity } from "../tools/notebook-mode/checkpoint-activity.ts";
import { resolveNotebookProject } from "../tools/notebook-mode/project-identity.ts";
import { scheduleNotebookStorageMaintenance } from "../tools/notebook-mode/storage-retention.ts";
import type { CodexExtensionRuntime } from "./runtime.ts";
import { discoverExternalNotesOwner } from "../context-management/external-notes-bridge.ts";
import type { CodexToolRegistration } from "./tools.ts";
import type { CodexUiController } from "./ui.ts";
import { registerCodexDeveloperMessageBroker } from "../developer-messages.ts";
import { registerCodexContextBriefingHost } from "../context-briefings.ts";
import type { ExtensionHandler, SessionStartEvent, ModelSelectEvent, SessionBeforeSwitchEvent, SessionBeforeForkEvent, SessionBeforeTreeEvent, SessionTreeEvent, SessionShutdownEvent, SessionBeforeSwitchResult, SessionBeforeForkResult, SessionBeforeTreeResult } from "@earendil-works/pi-coding-agent";

export function createCodexSessionLifecycle(
	pi: ExtensionAPI,
	runtime: CodexExtensionRuntime,
	tools: CodexToolRegistration,
	ui: CodexUiController,
	codeMode: CodeModeRegistration,
	proxyProvider: CodeModeProxyProviderRegistration,
	modelSelected: (ctx: ExtensionContext) => void,
) {
	const { state, tracker, sessions } = runtime;
	let activeContext: ExtensionContext | undefined;
	let checkpointActivity: NotebookCheckpointActivity | undefined;
	let checkpointSession: string | undefined;
	const trackCheckpoint = async (ctx: ExtensionContext): Promise<void> => {
		checkpointActivity?.release();
		checkpointSession = notebookSessionIdentity(ctx);
		checkpointActivity = await activateNotebookCheckpoint({
			project: resolveNotebookProject(ctx.cwd), session: checkpointSession, agentDir: getAgentDir(),
		});
		scheduleNotebookStorageMaintenance(getAgentDir());
	};
	let pendingExtensionToolRefresh = false;
	const unregisterContextBriefingHost = registerCodexContextBriefingHost(pi);
	const unregisterDeveloperMessageBroker = registerCodexDeveloperMessageBroker(
		pi,
		() => Boolean(
			activeContext &&
			supportsCodexDeveloperMessages(activeContext, state),
		),
	);
	const unregisterExtensionToolRefresh = onCodeModeExtensionToolsRefresh(
		pi,
		() => {
			if (!activeContext) return;
			if (!activeContext.isIdle()) {
				pendingExtensionToolRefresh = true;
				return;
			}
			pendingExtensionToolRefresh = false;
			syncAdapter(pi, activeContext, state);
		},
	);
	const activate = (ctx: ExtensionContext): void => {
		if (checkpointSession === notebookSessionIdentity(ctx)) checkpointActivity?.touch();
		activeContext = ctx;
		pendingExtensionToolRefresh = false;
	};
	const flushToolRefresh = (ctx: ExtensionContext): void => {
		if (pendingExtensionToolRefresh) {
			pendingExtensionToolRefresh = false;
			syncAdapter(pi, ctx, state);
		}
	};

	return {
		start: async (event, ctx) => {
			discoverExternalNotesOwner(pi, state);
			await trackCheckpoint(ctx);
			activate(ctx);
			state.notebookStatusMessageId = undefined;
			ui.invalidateUsageStatus();
			await runtime.lanVoice.stop(ctx);
			runtime.voice.resetContextAnnouncements();
			runtime.voice.resetSessionContext();
			initializeBashParser();
			runtime.resetTransport();
			state.developerMessages.clear();
			state.contextWindows.reset();
			state.contextKickoff.reset();
			if (!state.externalNotes) state.contextTree.beginSession(pi);
			runtime.backgroundWidget.ctx = ctx;
			state.cwd = ctx.cwd;
			state.config = readEffectiveCodexConversionConfig({
				cwd: ctx.cwd,
				projectTrusted: ctx.isProjectTrusted(),
			});
			state.executionMode = state.config.executionMode;
			proxyProvider.applyConfig(state.config, ctx.modelRegistry);
			state.promptSkills = extractPiPromptSkills(ctx.getSystemPrompt());
			if (state.config.voiceFeaturesOnly) {
				clearApplyPatchRenderState();
				ui.clearBackgroundWidget();
				syncAdapter(pi, ctx, state);
				await runtime.configureDiagnostics(ctx);
				return;
			}
			sessions.setBaseEnv(runtime.execEnv());
			tracker.clear();
			clearApplyPatchRenderState();
			ui.renderBackgroundWidget();
			syncAdapter(pi, ctx, state);
			// A fresh worker may adopt its family before its first prepared turn.
			if (!state.externalNotes) state.contextWindows.restore(ctx.sessionManager.getBranch());
			await runtime.configureDiagnostics(ctx);
			void ui.refreshUsageStatus(ctx);
			prepareCodeModeHost(codeMode, ctx);
			if (event.reason === "startup") await maybeWarnLocalCheckoutVersion(ctx);
		},
		modelSelected: async (_event, ctx) => {
			if (!state.externalNotes) state.contextTree.handoff.reset();
			modelSelected(ctx);
			activate(ctx);
			ui.invalidateUsageStatus();
			runtime.resetTransport(ctx.sessionManager.getSessionId());
			state.cwd = ctx.cwd;
			state.promptSkills = extractPiPromptSkills(ctx.getSystemPrompt());
			proxyProvider.applyConfig(state.config, ctx.modelRegistry);
			if (state.config.voiceFeaturesOnly) {
				ui.clearBackgroundWidget();
				syncAdapter(pi, ctx, state);
				await runtime.configureDiagnostics(ctx);
				return;
			}
			const plan = syncAdapter(pi, ctx, state);
			// Keep existing window state, but leave initial creation to the first prepared turn.
			if (!state.externalNotes) state.contextWindows.ensureInitialized(pi, ctx, plan.contextManagement, { createIfMissing: false });
			await runtime.configureDiagnostics(ctx);
			void ui.refreshUsageStatus(ctx);
			prepareCodeModeHost(codeMode, ctx);
		},
		beforeSwitch: () => !state.externalNotes && (state.contextTree.handoff.active || state.contextKickoff.hasIdleInput) ? { cancel: true } : undefined,
		beforeFork: () => !state.externalNotes && (state.contextTree.handoff.active || state.contextKickoff.hasIdleInput) ? { cancel: true } : undefined,
		beforeTree: (event, ctx) => {
			if (state.externalNotes) return;
			if (state.contextTree.handoff.active) return { cancel: true };
			if (state.contextTree.archiving) return;
			if (state.contextKickoff.hasIdleInput) return { cancel: true };
			const plan = resolveCodexRuntimePlanForState(ctx, state);
			if (!plan.notesTreeHandoff || !event.preparation.userWantsSummary) return;
			return state.contextTree.handoff.prepare(pi, event, ctx, plan.contextManagementMode);
		},
		tree: async (event, ctx) => {
			activate(ctx);
			const previousMode = state.executionMode;
			runtime.resetTransport(ctx.sessionManager.getSessionId());
			// Internal rollover preserves the admission that requested it.
			if (!state.externalNotes && state.contextTree.handleSessionTree(event)) return;
			state.notebookStatusMessageId = undefined;
			if (previousMode === "notebook" || state.executionMode === "notebook") appendNotebookTreeEpoch(pi);
			await codeMode.shutdownHost();
			await trackCheckpoint(ctx);
			proxyProvider.applyConfig(state.config, ctx.modelRegistry);
			const plan = syncAdapter(pi, ctx, state);
			if (!state.externalNotes) state.contextWindows.ensureInitialized(pi, ctx, plan.contextManagement, { createIfMissing: false });
			prepareCodeModeHost(codeMode, ctx);
			if (previousMode === "notebook" || state.executionMode === "notebook") {
				ctx.ui.notify("Notebook state reset after conversation-tree navigation", "info");
			}
		},
		shutdown: async (_event, ctx) => {
			const failures: unknown[] = [];
			pendingExtensionToolRefresh = false;
			await runShutdownStep(failures, unregisterExtensionToolRefresh);
			await runShutdownStep(failures, () => ui.invalidateBackgroundWidget());
			await runShutdownStep(failures, () => runtime.lanVoice.stop(ctx));
			await runShutdownStep(failures, () => runtime.voice.stop({ announce: true }));
			// Voice's persisted end policy still needs the active developer broker.
			activeContext = undefined;
			await runShutdownStep(failures, unregisterDeveloperMessageBroker);
			await runShutdownStep(failures, unregisterContextBriefingHost);
			await runShutdownStep(failures, () => runtime.shutdownTransport(ctx.sessionManager.getSessionId()));
			await runShutdownStep(failures, () => runtime.shutdownDiagnostics());
			await runShutdownStep(failures, () => sessions.shutdown());
			await runShutdownStep(failures, () => tools.shutdown());
			await runShutdownStep(failures, () => proxyProvider.shutdown());
			await runShutdownStep(failures, () => codeMode.shutdown());
			await runShutdownStep(failures, () => checkpointActivity?.release());
			state.developerMessages.clear();
			state.contextWindows.reset();
			state.contextKickoff.reset();
			state.contextTree.reset();
			if (failures.length === 1) throw failures[0];
			if (failures.length > 1) throw new AggregateError(failures, "Codex extension shutdown failed");
		},
		activate,
		flushToolRefresh,
	} satisfies {
		activate: (ctx: ExtensionContext) => void;
		flushToolRefresh: (ctx: ExtensionContext) => void;
		start: ExtensionHandler<SessionStartEvent>;
		modelSelected: ExtensionHandler<ModelSelectEvent>;
		beforeSwitch: ExtensionHandler<SessionBeforeSwitchEvent, SessionBeforeSwitchResult>;
		beforeFork: ExtensionHandler<SessionBeforeForkEvent, SessionBeforeForkResult>;
		beforeTree: ExtensionHandler<SessionBeforeTreeEvent, SessionBeforeTreeResult>;
		tree: ExtensionHandler<SessionTreeEvent>;
		shutdown: ExtensionHandler<SessionShutdownEvent>;
	};
}

function isAbortError(error: unknown): boolean {
	return error instanceof Error && (
		error.name === "AbortError"
		|| error.name === "ABORT_ERR"
		|| (error as Error & { code?: unknown }).code === "ABORT_ERR"
	);
}

export function prepareCodeModeHost(codeMode: CodeModeRegistration, ctx: ExtensionContext): void {
	void codeMode.prepare(ctx)?.catch((error: unknown) => {
		if (isAbortError(error)) return;
		ctx.ui.notify(`Code Mode host setup failed: ${error instanceof Error ? error.message : String(error)}`, "error");
	});
}

async function runShutdownStep(failures: unknown[], action: () => unknown): Promise<void> {
	try {
		await action();
	} catch (error) {
		failures.push(error);
	}
}