import {
	CONFIG_DIR_NAME,
	type ExtensionAPI,
	type ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import { Text } from "@earendil-works/pi-tui";
import {
	deriveNativeCompactionSetup,
	detectNativeCompaction,
	loadConfig,
	publicConfigFromResolved,
	validateConfigDraftForPath,
	writeConfigToDisk,
	writeProjectNativeCompactionSettings,
	type ContinuityConfig,
	type NativeCompactionStatus,
	type ResolvedContinuityConfig,
	type SynthesisEffort,
} from "../../src/config.js";
import { PRODUCT_NAME } from "../../src/constants.js";
import {
	completeNativeCompaction,
	createHandoffState,
	prepareNativeCompaction,
	runContinuityCheckpoint,
} from "../../src/handoff.js";
import {
	findStalePendingArtifacts,
	footerStatusForMessage,
	formatSettings,
	formatStatus,
	nativeViewForStatus,
	notifyStatus,
	statusHeadlineForConfig,
} from "../../src/status.js";

interface RuntimeContinuitySettingsState {
	liveNative?: NativeCompactionStatus;
	pendingNative?: NativeCompactionStatus;
	restartRequired: boolean;
}

interface SessionContinuityRegistry {
	activeRuntimeKeys: WeakSet<object>;
	settingsByRuntime: WeakMap<object, RuntimeContinuitySettingsState>;
}

const REGISTRY_SYMBOL = Symbol.for("pi-session-continuity.registry.v2");

function singletonRegistry(): SessionContinuityRegistry {
	const globalObject = globalThis as typeof globalThis & {
		[REGISTRY_SYMBOL]?: SessionContinuityRegistry;
	};
	return (globalObject[REGISTRY_SYMBOL] ??= {
		activeRuntimeKeys: new WeakSet<object>(),
		settingsByRuntime: new WeakMap<object, RuntimeContinuitySettingsState>(),
	});
}

function runtimeKeyFor(pi: ExtensionAPI): object {
	return (pi as unknown as { events?: object }).events ?? pi;
}

function modelId(ctx: ExtensionContext): string {
	return ctx.model
		? `${ctx.model.provider}/${ctx.model.id}`
		: "unknown/unknown";
}

export default function sessionContinuityExtension(pi: ExtensionAPI) {
	const registry = singletonRegistry();
	const runtimeKey = runtimeKeyFor(pi);
	if (registry.activeRuntimeKeys.has(runtimeKey)) {
		console.warn(
			"pi-session-continuity: duplicate package registration ignored; /continuity will be served by the already-registered copy.",
		);
		return;
	}
	registry.activeRuntimeKeys.add(runtimeKey);
	const runtimeSettings = registry.settingsByRuntime.get(runtimeKey) ?? {
		restartRequired: false,
	};
	registry.settingsByRuntime.set(runtimeKey, runtimeSettings);

	const state = createHandoffState();
	let currentConfig: ResolvedContinuityConfig | undefined;
	let staleNotificationSent = false;

	pi.registerMessageRenderer(
		"session-continuity-panel",
		(message, _options, theme) => {
			const content =
				typeof message.content === "string" ? message.content : "";
			const [headline = "", ...rest] = content.split("\n");
			return new Text([theme.bold(headline), ...rest].join("\n"), 1, 1);
		},
	);

	function showCommandOutput(ctx: ExtensionContext, message: string): void {
		if (!ctx.hasUI) {
			process.stdout.write(`${message}\n`);
			return;
		}
		notifyStatus(ctx, message);
	}

	function showCommandPanel(ctx: ExtensionContext, message: string): void {
		if (ctx.hasUI && ctx.mode === "tui") {
			pi.sendMessage({
				customType: "session-continuity-panel",
				content: message,
				display: true,
				details: { product: PRODUCT_NAME },
			});
			ctx.ui.setStatus("session-continuity", footerStatusForMessage(message));
			return;
		}
		showCommandOutput(ctx, message);
	}

	async function refreshConfig(
		ctx: ExtensionContext,
	): Promise<ResolvedContinuityConfig> {
		currentConfig = await loadConfig(ctx, CONFIG_DIR_NAME);
		return currentConfig;
	}

	function liveNative(ctx: ExtensionContext): NativeCompactionStatus {
		return (runtimeSettings.liveNative ??= detectNativeCompaction(
			ctx.cwd,
			CONFIG_DIR_NAME,
			ctx.isProjectTrusted(),
		));
	}

	function pendingNative(
		ctx: ExtensionContext,
	): NativeCompactionStatus | undefined {
		if (!runtimeSettings.restartRequired) return undefined;
		return (runtimeSettings.pendingNative ??= detectNativeCompaction(
			ctx.cwd,
			CONFIG_DIR_NAME,
			ctx.isProjectTrusted(),
		));
	}

	async function saveConfigDraft(
		ctx: ExtensionContext,
		draft: ContinuityConfig,
	): Promise<ResolvedContinuityConfig | undefined> {
		const config = currentConfig ?? (await refreshConfig(ctx));
		const errors = validateConfigDraftForPath(draft, config.configPath);
		if (errors.length > 0) {
			ctx.ui.notify(
				`${PRODUCT_NAME}: settings not saved. ${errors.join("; ")}`,
				"warning",
			);
			return undefined;
		}
		await writeConfigToDisk(config.configPath, draft);
		const refreshed = await refreshConfig(ctx);
		ctx.ui.notify(`${PRODUCT_NAME}: extension settings saved.`, "info");
		return refreshed;
	}

	async function configureNativeCompaction(
		ctx: ExtensionContext,
	): Promise<void> {
		const contextWindow = ctx.model?.contextWindow ?? 0;
		if (contextWindow <= 0) {
			ctx.ui.notify(
				`${PRODUCT_NAME}: active model context window is unavailable.`,
				"warning",
			);
			return;
		}
		const triggerChoice = await ctx.ui.select(
			"Approximate trigger for the current model",
			["50%", "55%", "60%", "65%", "70%", "75%", "80%", "85%", "90%", "95%"],
		);
		if (!triggerChoice) return;
		const triggerPercent = Number.parseInt(triggerChoice, 10);
		const keepChoice = await ctx.ui.select(
			"Approximate recent context to keep",
			["5%", "10%", "15%", "20%", "25%"].filter(
				(choice) => Number.parseInt(choice, 10) < triggerPercent,
			),
		);
		if (!keepChoice) return;
		const setup = deriveNativeCompactionSetup(
			contextWindow,
			triggerPercent,
			Number.parseInt(keepChoice, 10),
		);
		const settingsPath = liveNative(ctx).settingsPath;
		const approved = await ctx.ui.confirm(
			"Apply native Pi compaction settings?",
			[
				`Model window: ${setup.contextWindow.toLocaleString("en-US")}`,
				`Approximate trigger: ${setup.triggerTokens.toLocaleString("en-US")} (${setup.triggerPercent}%)`,
				`reserveTokens: ${setup.reserveTokens.toLocaleString("en-US")}`,
				`keepRecentTokens: ${setup.keepRecentTokens.toLocaleString("en-US")} (${setup.keepRecentPercent}%)`,
				`Write: ${settingsPath}`,
				"A full Pi restart is required.",
			].join("\n"),
		);
		if (!approved) return;
		const writtenPath = await writeProjectNativeCompactionSettings(
			ctx.cwd,
			CONFIG_DIR_NAME,
			setup,
		);
		runtimeSettings.pendingNative = detectNativeCompaction(
			ctx.cwd,
			CONFIG_DIR_NAME,
			ctx.isProjectTrusted(),
		);
		runtimeSettings.restartRequired = true;
		ctx.ui.notify(
			`${PRODUCT_NAME}: settings saved to ${writtenPath}; restart Pi to activate them.`,
			"warning",
		);
	}

	async function openSettingsMenu(
		ctx: ExtensionContext,
		initialConfig: ResolvedContinuityConfig,
	): Promise<void> {
		const native = liveNative(ctx);
		const view = nativeViewForStatus(native, ctx.model?.contextWindow ?? 0);
		const pending = pendingNative(ctx);
		const pendingView = pending
			? nativeViewForStatus(pending, ctx.model?.contextWindow ?? 0)
			: undefined;
		if (!ctx.hasUI || ctx.mode !== "tui" || !initialConfig.trusted) {
			showCommandPanel(
				ctx,
				formatSettings(initialConfig, view, native.settingsPath, pendingView),
			);
			return;
		}

		let config = initialConfig;
		for (;;) {
			const choice = await ctx.ui.select("Pi Session Continuity settings", [
				`Enabled: ${config.enabled ? "on" : "off"}`,
				`Native compaction setup: reserve ${native.reserveTokens.toLocaleString("en-US")} · keep ${native.keepRecentTokens.toLocaleString("en-US")}`,
				`Synthesis model: ${config.synthesisModel}`,
				`Synthesis effort: ${config.synthesisEffort}`,
				`Artifact directory: ${config.artifactDirectory}`,
				"Done",
			]);
			if (!choice || choice === "Done") return;
			if (choice.startsWith("Native compaction setup:")) {
				await configureNativeCompaction(ctx);
				continue;
			}

			const draft = publicConfigFromResolved(config);
			if (choice.startsWith("Enabled:")) {
				const enabled = await ctx.ui.select("Enable Session Continuity?", [
					"on",
					"off",
				]);
				if (!enabled) continue;
				draft.enabled = enabled === "on";
			} else if (choice.startsWith("Synthesis model:")) {
				const synthesisModel = await ctx.ui.input(
					'Synthesis model ("inherit" or provider/model)',
					config.synthesisModel,
				);
				if (!synthesisModel) continue;
				draft.synthesisModel = synthesisModel.trim();
			} else if (choice.startsWith("Synthesis effort:")) {
				const effort = await ctx.ui.select("Synthesis effort", [
					"inherit",
					"minimal",
					"low",
					"medium",
					"high",
					"xhigh",
				]);
				if (!effort) continue;
				draft.synthesisEffort = effort as SynthesisEffort;
			} else if (choice.startsWith("Artifact directory:")) {
				const artifactDirectory = await ctx.ui.input(
					"Artifact directory (relative to Pi config dir unless absolute)",
					config.artifactDirectory,
				);
				if (!artifactDirectory) continue;
				draft.artifactDirectory = artifactDirectory.trim();
			}
			const refreshed = await saveConfigDraft(ctx, draft);
			if (refreshed) config = refreshed;
		}
	}

	async function statusPanel(ctx: ExtensionContext): Promise<string> {
		const config = currentConfig ?? (await refreshConfig(ctx));
		const native = liveNative(ctx);
		const view = nativeViewForStatus(native, ctx.model?.contextWindow ?? 0);
		const pending = pendingNative(ctx);
		const pendingView = pending
			? nativeViewForStatus(pending, ctx.model?.contextWindow ?? 0)
			: undefined;
		const stale = config.trusted
			? await findStalePendingArtifacts(
					config.artifactDirectoryPath,
					ctx.sessionManager.getSessionId(),
				)
			: [];
		return formatStatus(
			config,
			view,
			native.settingsPath,
			state,
			modelId(ctx),
			stale[0]?.path,
			runtimeSettings.restartRequired,
			pendingView,
		);
	}

	pi.on("session_start", async (_event, ctx) => {
		const config = await refreshConfig(ctx);
		const native = liveNative(ctx);
		const view = nativeViewForStatus(native, ctx.model?.contextWindow ?? 0);
		if (config.trusted) {
			const stale = await findStalePendingArtifacts(
				config.artifactDirectoryPath,
				ctx.sessionManager.getSessionId(),
			);
			if (stale[0] && !staleNotificationSent) {
				staleNotificationSent = true;
				ctx.ui.notify(
					`${PRODUCT_NAME}: stale pending Continuity Brief found for this session; leaving it inert. ${stale[0].path}`,
					"warning",
				);
			}
		}
		notifyStatus(ctx, statusHeadlineForConfig(config, view));
	});

	pi.on("session_before_compact", async (event, ctx) => {
		const config = await refreshConfig(ctx);
		if (!config.trusted || !config.enabled) return undefined;
		if (!config.valid) {
			ctx.ui.notify(
				`${PRODUCT_NAME} failed: invalid configuration in ${config.configPath}. Native compaction was cancelled.`,
				"error",
			);
			return { cancel: true };
		}
		return prepareNativeCompaction(ctx, config, state, event);
	});

	pi.on("session_compact", async (event, ctx) => {
		const config = currentConfig ?? (await refreshConfig(ctx));
		await completeNativeCompaction(pi, ctx, config, state, event);
	});

	pi.registerCommand("continuity", {
		description: "Pi Session Continuity: status, checkpoint, and settings",
		getArgumentCompletions(prefix) {
			return ["status", "checkpoint", "settings"]
				.filter((item) => item.startsWith(prefix))
				.map((value) => ({ value, label: value }));
		},
		handler: async (args, ctx) => {
			const subcommand = args.trim().split(/\s+/).filter(Boolean)[0] ?? "";
			const config = await refreshConfig(ctx);
			const native = liveNative(ctx);

			if (subcommand === "") {
				if (ctx.hasUI && ctx.mode === "tui" && config.trusted) {
					const choice = await ctx.ui.select("Pi Session Continuity", [
						"Status",
						"Create checkpoint now",
						"Settings",
						"Done",
					]);
					if (!choice || choice === "Done") return;
					if (choice === "Status") {
						showCommandPanel(ctx, await statusPanel(ctx));
						return;
					}
					if (choice === "Settings") {
						await openSettingsMenu(ctx, config);
						return;
					}
					if (choice === "Create checkpoint now") {
						await runContinuityCheckpoint(ctx, config, state, native);
						return;
					}
				}
				showCommandPanel(ctx, await statusPanel(ctx));
				return;
			}

			if (subcommand === "status") {
				showCommandPanel(ctx, await statusPanel(ctx));
				return;
			}
			if (subcommand === "settings") {
				await openSettingsMenu(ctx, config);
				return;
			}
			if (subcommand === "checkpoint") {
				if (!config.trusted || !config.valid || !config.enabled) {
					showCommandOutput(
						ctx,
						`${PRODUCT_NAME}: checkpoint disabled. ${config.disabledReason ?? config.errors.join("; ")}`,
					);
					return;
				}
				await runContinuityCheckpoint(ctx, config, state, native);
				return;
			}
			showCommandOutput(
				ctx,
				`${PRODUCT_NAME}: unknown subcommand '${subcommand}'. Use /continuity status, /continuity checkpoint, or /continuity settings.`,
			);
		},
	});

	pi.on("session_shutdown", () => {
		registry.activeRuntimeKeys.delete(runtimeKey);
	});
}