import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
	COMMAND_NAME,
	COMPACTION_CONTEXT_TYPE,
	DEFAULT_CONFIG,
	EXTENSION_NAME,
	LEGACY_CONFIG_PATH,
	PROJECT_CONTEXT_TYPE,
} from "./constants.js";
import { ensureJitiFsCacheDirectory } from "./jiti-cache.js";
import type { ContextInjectorConfig } from "./types.js";

type ConfigStoreModule = typeof import("./config-store.js");
type ContextBuilderModule = typeof import("./context-builder.js");
type CompactionDedupeModule = typeof import("./compaction-dedupe.js");
type ConfigModalModule = typeof import("./config-modal.js");
type LoggerModule = typeof import("./logger.js");

interface ContextInjectorLoggerLike {
	debug(message: string, details?: unknown): void;
	warn(message: string, details?: unknown): void;
}

function cloneDefaultConfig(): ContextInjectorConfig {
	return {
		...DEFAULT_CONFIG,
		ignoredSections: [...DEFAULT_CONFIG.ignoredSections],
		compaction: {
			...DEFAULT_CONFIG.compaction,
			additionalContext: [...DEFAULT_CONFIG.compaction.additionalContext],
		},
	};
}

function getSessionKey(ctx: ExtensionContext): string {
	return ctx.sessionManager.getSessionFile() ?? ctx.sessionManager.getSessionId();
}

function sessionAlreadyHasInjectedProjectContext(ctx: ExtensionContext): boolean {
	return ctx.sessionManager.getEntries().some((entry) => {
		if (entry.type !== "message") {
			return false;
		}
		const message = entry.message as { role?: string; customType?: string };
		return message.role === "custom" && message.customType === PROJECT_CONTEXT_TYPE;
	});
}

function sessionAlreadyHasAssistantReply(ctx: ExtensionContext): boolean {
	return ctx.sessionManager.getEntries().some((entry) => {
		if (entry.type !== "message") {
			return false;
		}
		const message = entry.message as { role?: string };
		return message.role === "assistant";
	});
}

function isParentLinkedSession(ctx: ExtensionContext): boolean {
	const header = ctx.sessionManager.getHeader();
	return Boolean(header?.parentSession);
}

function shouldSkipParentLinkedSessionContext(config: ContextInjectorConfig, ctx: ExtensionContext): boolean {
	return config.skipForkedSessions && isParentLinkedSession(ctx);
}

export default function contextInjectorExtension(pi: ExtensionAPI): void {
	let config: ContextInjectorConfig = cloneDefaultConfig();
	let pendingLoadWarning: string | undefined;
	let configStorePromise: Promise<ConfigStoreModule> | undefined;
	let contextBuilderPromise: Promise<ContextBuilderModule> | undefined;
	let compactionDedupePromise: Promise<CompactionDedupeModule> | undefined;
	let configModalPromise: Promise<ConfigModalModule> | undefined;
	let loggerModulePromise: Promise<LoggerModule> | undefined;
	let loggerInstance: InstanceType<LoggerModule["ContextInjectorLogger"]> | undefined;
	const warnedMessages = new Set<string>();
	// before_agent_start runs for every prompt, but project context should only be
	// decided once per session: inject on the first eligible turn, then never
	// reinject on later turns in that same session.
	const initialProjectContextHandledSessions = new Set<string>();

	const importRuntimeModule = <T>(loader: () => Promise<T>): Promise<T> => {
		ensureJitiFsCacheDirectory();
		return loader();
	};

	const loadConfigStore = (): Promise<ConfigStoreModule> => {
		configStorePromise ??= importRuntimeModule(() => import("./config-store.js"));
		return configStorePromise;
	};

	const loadContextBuilder = (): Promise<ContextBuilderModule> => {
		contextBuilderPromise ??= importRuntimeModule(() => import("./context-builder.js"));
		return contextBuilderPromise;
	};

	const loadCompactionDedupe = (): Promise<CompactionDedupeModule> => {
		compactionDedupePromise ??= importRuntimeModule(() => import("./compaction-dedupe.js"));
		return compactionDedupePromise;
	};

	const loadConfigModal = (): Promise<ConfigModalModule> => {
		configModalPromise ??= importRuntimeModule(() => import("./config-modal.js"));
		return configModalPromise;
	};

	const writeLog = (level: "debug" | "warn", message: string, details?: unknown): void => {
		if (!config.debug) {
			return;
		}

		loggerModulePromise ??= importRuntimeModule(() => import("./logger.js"));
		void loggerModulePromise.then((moduleValue) => {
			loggerInstance ??= new moduleValue.ContextInjectorLogger(() => config.debug);
			loggerInstance[level](message, details);
		}).catch(() => undefined);
	};

	const logger: ContextInjectorLoggerLike = {
		debug: (message, details) => writeLog("debug", message, details),
		warn: (message, details) => writeLog("warn", message, details),
	};

	const warnOnce = (message: string, ctx?: Pick<ExtensionContext, "hasUI" | "ui">): void => {
		if (warnedMessages.has(message)) {
			return;
		}

		warnedMessages.add(message);
		logger.warn(message);
		if (ctx?.hasUI) {
			ctx.ui.notify(message, "warning");
		}
	};

	const refreshConfig = async (
		ctx?: Pick<ExtensionContext, "hasUI" | "ui">,
		options: { ensureExists?: boolean } = {},
	): Promise<ConfigStoreModule> => {
		const configStore = await loadConfigStore();
		if (options.ensureExists) {
			const ensureResult = configStore.ensureConfigExists();
			if (ensureResult.error) {
				warnOnce(ensureResult.error, ctx);
			}
		}

		const loaded = configStore.loadContextInjectorConfig();
		config = loaded.config;
		pendingLoadWarning = loaded.warning;

		if (loaded.source === "legacy") {
			warnOnce(
				`${EXTENSION_NAME}: using legacy config ${LEGACY_CONFIG_PATH}. Create ${configStore.getContextInjectorConfigPath()} to override it.`,
				ctx,
			);
		}

		if (pendingLoadWarning) {
			warnOnce(pendingLoadWarning, ctx);
			pendingLoadWarning = undefined;
		}

		return configStore;
	};

	pi.registerCommand(COMMAND_NAME, {
		description: "Configure project context injection",
		handler: async (args, ctx) => {
			const [configStore, configModal] = await Promise.all([
				refreshConfig(ctx, { ensureExists: true }),
				loadConfigModal(),
			]);

			await configModal.handleContextInjectorCommand(args, ctx, {
				getConfig: () => config,
				setConfig: (next: ContextInjectorConfig, commandCtx: ExtensionCommandContext) => {
					config = configStore.normalizeContextInjectorConfig(next);
					const saved = configStore.saveContextInjectorConfig(config);
					if (!saved.success && saved.error) {
						commandCtx.ui.notify(saved.error, "error");
					}
				},
				getConfigPath: configStore.getContextInjectorConfigPath,
			});
		},
	});

	pi.on("before_agent_start", async (event, ctx) => {
		await refreshConfig(ctx);
		if (!config.enabled) {
			return {};
		}

		const sessionKey = getSessionKey(ctx);
		if (initialProjectContextHandledSessions.has(sessionKey)) {
			return {};
		}

		if (shouldSkipParentLinkedSessionContext(config, ctx)) {
			initialProjectContextHandledSessions.add(sessionKey);
			logger.debug("Skipped initial project context for parent-linked session", { sessionKey });
			return {};
		}

		if (sessionAlreadyHasInjectedProjectContext(ctx) || sessionAlreadyHasAssistantReply(ctx)) {
			initialProjectContextHandledSessions.add(sessionKey);
			return {};
		}

		try {
			const { buildProjectContext, detectFormat } = await loadContextBuilder();
			const format = detectFormat(config, ctx.model);
			const built = await buildProjectContext(ctx.cwd, format, config, logger);
			if (!built.block) {
				initialProjectContextHandledSessions.add(sessionKey);
				logger.debug("No project context sources available for initial session injection.", { sessionKey });
				return {};
			}

			initialProjectContextHandledSessions.add(sessionKey);
			logger.debug("Injecting initial project context", {
				sessionKey,
				format,
				sections: built.sectionNames,
				target: config.injectionTarget,
			});

			if (config.injectionTarget === "system_prompt") {
				return {
					systemPrompt: `${event.systemPrompt}\n\n${built.block}`,
				};
			}

			return {
				message: {
					customType: PROJECT_CONTEXT_TYPE,
					content: built.block,
					display: !config.silent,
					details: {
						format,
						sections: built.sectionNames,
						generatedAt: new Date().toISOString(),
					},
				},
			};
		} catch (error) {
			const message = error instanceof Error ? error.message : String(error);
			warnOnce(`${EXTENSION_NAME}: failed to inject project context: ${message}`, ctx);
			return {};
		}
	});

	pi.on("session_compact", async (event, ctx) => {
		await refreshConfig(ctx);
		if (!config.enabled || !config.compaction.enabled) {
			return;
		}

		if (shouldSkipParentLinkedSessionContext(config, ctx)) {
			return;
		}

		try {
			const { buildCompactionContext, detectFormat, extractTodoSnapshotFromBranch } = await loadContextBuilder();
			const format = detectFormat(config, ctx.model);
			const todoSnapshot = extractTodoSnapshotFromBranch(ctx.sessionManager.getBranch() as unknown[]);
			const built = await buildCompactionContext(ctx.cwd, format, config, logger, todoSnapshot);
			if (!built.block) {
				logger.debug("Compaction context generation skipped (no sections).", { session: getSessionKey(ctx) });
				return;
			}

			const sessionKey = getSessionKey(ctx);
			const compactionEntryId = typeof event.compactionEntry.id === "string" && event.compactionEntry.id.trim()
				? event.compactionEntry.id.trim()
				: undefined;
			const {
				createCompactionContextHash,
				createCompactionContextMetadata,
				sessionAlreadyHasCompactionContext,
			} = await loadCompactionDedupe();
			const contextHash = createCompactionContextHash(built.block);
			if (
				sessionAlreadyHasCompactionContext(ctx.sessionManager.getEntries() as unknown[], {
					compactionEntryId,
					contextHash,
				})
			) {
				logger.debug("Compaction context already persisted; skipping reinjection.", {
					sessionKey,
					compactionEntryId,
				});
				return;
			}

			pi.sendMessage(
				{
					customType: COMPACTION_CONTEXT_TYPE,
					content: built.block,
					display: !config.silent,
					details: {
						...createCompactionContextMetadata({ compactionEntryId, contextHash }),
						format,
						sections: built.sectionNames,
						generatedAt: new Date().toISOString(),
					},
				},
				{ triggerTurn: false },
			);
		} catch (error) {
			const message = error instanceof Error ? error.message : String(error);
			warnOnce(`${EXTENSION_NAME}: failed to inject compaction context: ${message}`, ctx);
		}
	});
}