import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import { loadConfig, setDebugLog, dbg, type CtxVarsConfig } from "./config.ts";
import { VarStore } from "./store.ts";
import { runCompaction } from "./compaction.ts";
import { populateFromMessage } from "./populate.ts";
import { findEntryId, derivedId } from "./ids.ts";
import { registerTools } from "./tools.ts";
import { join } from "node:path";

/**
 * Context Variables — variable-based context management for pi.
 *
 * Replaces monolithic compaction with an agent-controlled variable store:
 * - Every summarized message becomes a variable (content + summary) in a SQLite store.
 * - The compaction agent decides pin / archive / drop per entry, writes per-entry
 *   summaries at archive time, and rolls old archives up when they accumulate.
 * - The main agent manages retention with context_pin / context_unpin /
 *   context_drop / context_archive / context_read / context_query.
 * - Nothing is destroyed: dropped and archived content stays fully recoverable.
 */

export default function (pi: ExtensionAPI) {
	const cfg: CtxVarsConfig = loadConfig();
	setDebugLog(cfg.debugLogPath);
	if (!cfg.enabled) {
		dbg("ctxvars disabled");
		return;
	}

	let currentStore: VarStore | null = null;

	function storeFor(ctx: ExtensionContext): VarStore | null {
		if (currentStore) return currentStore;
		try {
			const sessionFile = ctx.sessionManager.getSessionFile();
			let path: string;
			if (cfg.storeDir) {
				const id = ctx.sessionManager.getSessionId();
				path = join(cfg.storeDir, `${id}.ctxvars.sqlite`);
			} else if (sessionFile) {
				path = `${sessionFile}.ctxvars.sqlite`;
			} else {
				const id = ctx.sessionManager.getSessionId();
				path = join(process.env.TMPDIR ?? "/tmp", `pi-ctxvars-${id}.sqlite`);
			}
			currentStore = new VarStore(path);
			dbg(`store opened: ${path}`);
			return currentStore;
		} catch (err) {
			dbg(`store open failed: ${err instanceof Error ? err.message : String(err)}`);
			return null;
		}
	}

	pi.on("session_start", async (_event, ctx) => {
		storeFor(ctx);
	});

	// Realtime population: every message is stored as a variable as it happens.
	// Context management (summaries, pin/archive/drop, dependencies) is still
	// done at compaction time by the compaction agent.
	pi.on("message_end", async (event, ctx) => {
		const store = currentStore;
		if (!store) return;
		const message = event.message as never;
		try {
			// Eager attempt: the entry usually does not exist yet (pi appends it
			// AFTER extension handlers run), so this normally falls back to the
			// deterministic derived id. The deferred pass below fixes it up.
			const entryId = findEntryId(ctx.sessionManager as never, message);
			populateFromMessage(store, cfg, message, entryId);
		} catch (err) {
			dbg(`populate error: ${err instanceof Error ? err.message : String(err)}`);
		}
		// Deferred pass: pi appends the entry synchronously right after the
		// handler returns, so a macrotask later the real entry id is resolvable.
		// Capture the manager now: accessing ctx itself after print-mode shutdown
		// or session replacement raises pi's stale-context guard.
		const storeAtEvent = store;
		const sessionManagerAtEvent = ctx.sessionManager;
		setTimeout(() => {
			try {
				// The session may have shut down or switched before this task runs.
				if (currentStore !== storeAtEvent) return;
				const entryId = findEntryId(sessionManagerAtEvent as never, message);
				if (!entryId) return;
				const varRow = storeAtEvent.getByEntry(derivedId(message));
				if (varRow) {
					// Re-key to the real entry id (keeps seq and state).
					storeAtEvent.updateEntryId(varRow.id, entryId);
				} else if (!storeAtEvent.getByEntry(entryId)) {
					populateFromMessage(storeAtEvent, cfg, message, entryId);
				}
			} catch (err) {
				dbg(`populate deferred error: ${err instanceof Error ? err.message : String(err)}`);
			}
		}, 0);
	});

	pi.on("session_before_compact", async (event, ctx) => {
		const store = storeFor(ctx);
		if (!store) return undefined;
		const result = await runCompaction(event as never, ctx, cfg, store);
		if (!result) return undefined;
		return {
			compaction: {
				summary: result.summary,
				firstKeptEntryId: event.preparation.firstKeptEntryId,
				tokensBefore: event.preparation.tokensBefore,
				usage: result.usage as never,
				details: result.details,
			},
		};
	});

	pi.on("session_shutdown", () => {
		if (currentStore) {
			currentStore.close();
			currentStore = null;
		}
	});

	registerTools(pi, cfg, () => currentStore);

	pi.registerCommand("ctxv", {
		description: "Show context-variables store stats",
		handler: async (_args, ctx) => {
			const store = storeFor(ctx);
			if (!store) {
				ctx.ui.notify("Store unavailable", "error");
				return;
			}
			ctx.ui.notify(JSON.stringify(store.stats(), null, 2), "info");
		},
	});

	pi.registerCommand("ctxv-reset", {
		description: "Reset the context-variables store for this session",
		handler: async (_args, ctx) => {
			const store = storeFor(ctx);
			if (!store) return;
			store.reset();
			ctx.ui.notify("Store reset", "info");
		},
	});
}