import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
import { uuidv7 } from "@earendil-works/pi-ai";
import type { CtxVarsConfig } from "./config.ts";
import { dbg } from "./config.ts";
import { VarStore, type PendingDecision, type VariableRow } from "./store.ts";
import { fullText, kindOf, estimateTokens, truncate, sampleLongText } from "./serialize.ts";
import { findEntryId, derivedId } from "./ids.ts";
import {
	buildCompactionMessages,
	parseCompactionOutput,
	type CandidateInfo,
	type CompactionPacket,
	type ContextSnapshot,
	type SnapshotItem,
	type SnapshotTopic,
} from "./prompts.ts";

interface SessionBeforeCompactEventLike {
	preparation: {
		messagesToSummarize: unknown[];
		turnPrefixMessages: unknown[];
		previousSummary: string | null;
		firstKeptEntryId: string;
		tokensBefore: number;
	};
	signal?: AbortSignal;
}

export interface CompactionResult {
	summary: string;
	usage?: unknown;
	details: Record<string, unknown>;
}

function chooseModel(ctx: ExtensionContext, cfg: CtxVarsConfig) {
	let model = ctx.model;
	if (cfg.compactionModel) {
		const slash = cfg.compactionModel.indexOf("/");
		const provider = slash === -1 ? null : cfg.compactionModel.slice(0, slash);
		const id = slash === -1 ? cfg.compactionModel : cfg.compactionModel.slice(slash + 1);
		if (provider) model = ctx.modelRegistry.find(provider, id) ?? undefined;
	}
	return model;
}

async function callCompactionAgent(ctx: ExtensionContext, cfg: CtxVarsConfig, packet: CompactionPacket, signal?: AbortSignal) {
	const model = chooseModel(ctx, cfg);
	if (!model) return null;
	const response = await ctx.modelRegistry.complete(
		model,
		{ messages: buildCompactionMessages(packet, cfg) },
		{
			maxTokens: cfg.compactionMaxTokens,
			signal,
			cacheRetention: "none",
			sessionId: uuidv7(),
		},
	);
	const text = response.content
		.filter((c): c is { type: "text"; text: string } => c.type === "text")
		.map((c) => c.text)
		.join("\n");
	return { text, usage: response.usage };
}

async function callRollupAgent(ctx: ExtensionContext, cfg: CtxVarsConfig, targets: VariableRow[], signal?: AbortSignal): Promise<string | null> {
	const model = chooseModel(ctx, cfg);
	if (!model || targets.length === 0) return null;
	const source = targets.map((v) => `- [var_${v.id}] ${v.summary || truncate(v.content, 300)}`).join("\n");
	const prompt = `Compress these older archived memory summaries into one dense historical summary. Preserve decisions, requirements, exact values, important paths, failures that affect future work, and unresolved items. Remove procedural narration. Reply with ONLY JSON: {"summary":"..."}\n\n${source}`;
	try {
		const response = await ctx.modelRegistry.complete(
			model,
			{ messages: [{ role: "user" as const, content: [{ type: "text" as const, text: prompt }], timestamp: Date.now() }] },
			{
				maxTokens: Math.min(1200, cfg.compactionMaxTokens),
				signal,
				cacheRetention: "none",
				sessionId: uuidv7(),
			},
		);
		const text = response.content
			.filter((c): c is { type: "text"; text: string } => c.type === "text")
			.map((c) => c.text)
			.join("\n");
		const start = text.indexOf("{");
		const end = text.lastIndexOf("}");
		if (start === -1 || end === -1) return null;
		const parsed = JSON.parse(text.slice(start, end + 1)) as { summary?: unknown };
		return typeof parsed.summary === "string" && parsed.summary.trim() ? parsed.summary.trim() : null;
	} catch (err) {
		dbg(`rollup follow-up failed: ${err instanceof Error ? err.message : String(err)}`);
		return null;
	}
}

function addPending(map: Map<string | number, PendingDecision[]>, key: string | number, decision: PendingDecision) {
	const rows = map.get(key) ?? [];
	rows.push(decision);
	map.set(key, rows);
}

function mergedPending(...groups: Array<PendingDecision[] | undefined>): PendingDecision[] {
	return groups.flatMap((g) => g ?? []).sort((a, b) => a.id - b.id);
}

function decisionLabel(d: PendingDecision): string {
	if (d.action === "unpin") return "unpin; then choose archive or drop";
	return `${d.action}${d.decay ? "(decay=true)" : "(decay=false)"}${d.summary ? `: ${d.summary}` : ""}`;
}

function persistRollup(store: VarStore, targets: VariableRow[], summary: string, cfg: CtxVarsConfig): number | null {
	if (targets.length === 0 || !summary.trim()) return null;
	const first = `var_${targets[0].id}`;
	const last = `var_${targets[targets.length - 1].id}`;
	const range = `${first}..${last}`;
	const rollupId = store.addRollup(range, clipAtWord(summary.trim(), cfg.rollupMaxChars));
	for (const v of targets) store.setRollup(v.id, rollupId);
	return rollupId;
}

export async function runCompaction(
	event: SessionBeforeCompactEventLike,
	ctx: ExtensionContext,
	cfg: CtxVarsConfig,
	store: VarStore,
): Promise<CompactionResult | null> {
	const { preparation, signal } = event;
	const all = [...(preparation.messagesToSummarize ?? []), ...(preparation.turnPrefixMessages ?? [])];
	if (all.length === 0) {
		dbg("compaction: no candidates, deferring to default");
		return null;
	}
	const started = Date.now();
	dbg(`compaction: start, ${all.length} messages`);

	try {
		// ---- normalize pending main-agent decisions ----
		const pending = store.pendingDecisions();
		const pendingByEntry = new Map<string | number, PendingDecision[]>();
		const pendingByVarId = new Map<string | number, PendingDecision[]>();
		for (const d of pending) {
			if (d.target_type === "entry") {
				addPending(pendingByEntry, d.target, d);
			} else {
				const v = store.getVarByRef(d.target);
				if (v) addPending(pendingByVarId, v.id, d);
			}
		}

		// ---- candidates ----
		const fullTextByEntry = new Map<string, string>();
		const existingByEntry = new Map<string, VariableRow>();
		const rowEntryIdByCandidate = new Map<string, string>();
		const pendingByCandidate = new Map<string, PendingDecision[]>();
		const candidateVarIds = new Set<number>();
		const candidates: CandidateInfo[] = [];
		const budgetChars = cfg.packetTokenBudget * 4;

		for (const m of all as Array<{ id?: string }>) {
			const fallbackId = derivedId(m as never);
			const realId = findEntryId(ctx.sessionManager as never, m as never);
			const id = realId ?? fallbackId;
			const existingReal = realId ? store.getByEntry(realId) : null;
			const existingFallback = store.getByEntry(fallbackId);
			const existing = existingReal ?? existingFallback;
			const rowId: string = existingReal && realId ? realId : existingFallback ? fallbackId : id;
			const decisionRows = mergedPending(pendingByEntry.get(id), existing ? pendingByVarId.get(existing.id) : undefined);
			const latest = decisionRows.at(-1);
			const releasesPersistentPin = !!latest && latest.action !== "pin";

			// Existing persistent pins are rendered directly from the store. They
			// only become candidates when a later main-agent action releases them.
			if (existing?.pinned === 2 && !releasesPersistentPin) continue;

			const text = fullText(m as never);
			const kind = kindOf(m as never);
			const bindingIsDecayablePin = latest?.action === "pin" && latest.decay === 1;
			rowEntryIdByCandidate.set(id, rowId);
			fullTextByEntry.set(id, text);
			if (existing) {
				existingByEntry.set(id, existing);
				candidateVarIds.add(existing.id);
			}
			if (decisionRows.length) pendingByCandidate.set(id, decisionRows);
			candidates.push({
				id,
				varRef: existing ? `var_${existing.id}` : undefined,
				kind,
				sizeTokens: estimateTokens(text),
				text: kind === "tool_result" ? sampleLongText(text, cfg.toolResultTruncateChars) : truncate(text, cfg.candidateTruncateChars),
				bindingDecision: latest && !bindingIsDecayablePin ? decisionLabel(latest) : undefined,
				previouslyPinned: existing?.pinned === 1 || bindingIsDecayablePin,
			});
		}

		// Variable decisions can target entries whose raw message is not in this
		// compaction span. Archive/unpin needs model judgment, so add a synthetic
		// candidate backed by the stored full content.
		for (const [varIdKey, rows] of pendingByVarId) {
			const varId = Number(varIdKey);
			if (candidateVarIds.has(varId)) continue;
			const v = store.getVar(varId);
			if (!v) continue;
			const latest = rows.at(-1);
			if (!latest || (latest.action !== "archive" && latest.action !== "unpin")) continue;
			const id = v.entry_id ?? `var_${v.id}`;
			rowEntryIdByCandidate.set(id, id);
			fullTextByEntry.set(id, v.content);
			existingByEntry.set(id, v);
			pendingByCandidate.set(id, rows);
			candidateVarIds.add(v.id);
			candidates.push({
				id,
				varRef: `var_${v.id}`,
				kind: v.kind,
				sizeTokens: v.size_tokens,
				text: v.kind === "tool_result" ? sampleLongText(v.content, cfg.toolResultTruncateChars) : truncate(v.content, cfg.candidateTruncateChars),
				bindingDecision: decisionLabel(latest),
				previouslyPinned: v.pinned > 0,
			});
		}

		// Packet budget: trim oldest candidate bodies first.
		let totalChars = candidates.reduce((sum, c) => sum + c.text.length + 100, 0);
		if (totalChars > budgetChars) {
			for (const c of candidates) {
				if (totalChars <= budgetChars) break;
				if (c.text.length > 200) {
					const saved = c.text.length - 200;
					c.text = truncate(c.text, 200);
					totalChars -= saved;
				}
			}
		}

		// Project the current pass as archives. This avoids waiting one extra
		// compaction before requesting a rollup.
		const active = store.listInContext().filter((v) => v.pinned === 0 && !candidateVarIds.has(v.id));
		const projectedExcess = Math.max(0, active.length + candidates.length - cfg.archiveKeepIndividual);
		const initialRollupRows = active.slice(0, Math.min(projectedExcess, active.length));
		let activeArchives = active
			.filter((v) => !initialRollupRows.some((r) => r.id === v.id))
			.map((v) => ({ ref: `var_${v.id}`, summary: v.summary }));
		if (totalChars + activeArchives.reduce((s, a) => s + a.summary.length + 20, 0) > budgetChars) {
			activeArchives = activeArchives.slice(-Math.max(15, cfg.archiveKeepIndividual));
		}

		const packet: CompactionPacket = {
			candidates,
			activeArchives,
			rollupTargets: initialRollupRows.map((v) => ({ ref: `var_${v.id}`, summary: v.summary })),
			previousSummary: preparation.previousSummary,
		};

		// ---- compaction agent call ----
		dbg(`compaction: calling agent (${candidates.length} candidates, ${activeArchives.length} active archives)`);
		const out = await callCompactionAgent(ctx, cfg, packet, signal);
		if (!out) return null;
		const parsed = parseCompactionOutput(out.text);
		if (!parsed) {
			dbg(`compaction: unparseable output (${out.text.slice(0, 200)}), deferring to default`);
			return null;
		}
		dbg(`compaction: parsed ${parsed.decisions.length} decisions, ${parsed.snapshot.topics.length} topics`);

		// ---- apply candidate decisions ----
		const decisionByEntry = new Map(parsed.decisions.map((d) => [d.id, d]));
		const consumedDecisionIds = new Set<number>();
		const processedVarIds = new Set<number>();
		let archivedCount = 0;

		for (const c of candidates) {
			const prior = existingByEntry.get(c.id);
			const rowId = rowEntryIdByCandidate.get(c.id) ?? c.id;
			const seq = prior ? prior.seq : store.nextSeq();
			const pendingRows = pendingByCandidate.get(c.id) ?? [];
			const binding = pendingRows.at(-1);
			const modelDecision = decisionByEntry.get(c.id);
			const bindingIsDecayablePin = binding?.action === "pin" && binding.decay === 1;
			let action: "pin" | "archive" | "drop";
			let summary: string | undefined;
			let dependsOn: string[] | undefined = modelDecision?.depends_on;
			// The compaction model may create only decayable pins. Persistent pins
			// are reserved for an explicit binding main-agent decision.
			let pinDecay = true;

			if (binding && !bindingIsDecayablePin) {
				if (binding.action === "unpin") {
					action = modelDecision?.action === "drop" ? "drop" : "archive";
					summary = modelDecision?.summary;
				} else if (binding.action === "pin") {
					action = "pin";
					summary = binding.summary ?? modelDecision?.summary;
					pinDecay = binding.decay === 1;
				} else if (binding.action === "drop") {
					action = "drop";
					summary = undefined;
				} else {
					action = "archive";
					summary = binding.summary ?? modelDecision?.summary;
				}
			} else {
				action = modelDecision?.action ?? "archive";
				summary = modelDecision?.summary;
			}

			const content = fullTextByEntry.get(c.id) ?? c.text;
			let varId: number;
			if (action === "pin") {
				varId = store.upsertByEntry(rowId, {
					kind: c.kind,
					content,
					summary: summary?.slice(0, cfg.summaryMaxChars) ?? prior?.summary ?? "",
					sizeTokens: c.sizeTokens,
					seq,
					pinned: pinDecay ? 1 : 2,
					dropped: 0,
					inContext: 1,
				});
			} else if (action === "archive") {
				varId = store.upsertByEntry(rowId, {
					kind: c.kind,
					content,
					summary: (summary ?? prior?.summary ?? "").slice(0, cfg.summaryMaxChars),
					sizeTokens: c.sizeTokens,
					seq,
					pinned: 0,
					dropped: 0,
					inContext: 1,
				});
				archivedCount++;
			} else {
				varId = store.upsertByEntry(rowId, {
					kind: c.kind,
					content,
					summary: "",
					sizeTokens: c.sizeTokens,
					seq,
					pinned: 0,
					dropped: 1,
					inContext: 0,
				});
			}
			processedVarIds.add(varId);
			if (dependsOn) addDeps(store, varId, dependsOn);
			for (const d of pendingRows) consumedDecisionIds.add(d.id);
		}

		// Pin/drop actions on existing var_N targets do not need a synthetic LLM
		// candidate, but their state and bookkeeping still apply at this boundary.
		for (const [varIdKey, rows] of pendingByVarId) {
			const varId = Number(varIdKey);
			if (processedVarIds.has(varId)) continue;
			const v = store.getVar(varId);
			const latest = rows.at(-1);
			if (!v || !latest) continue;
			if (latest.action === "pin") {
				store.setPinned(varId, latest.decay ? 1 : 2);
				store.setDropped(varId, 0, 1);
			} else if (latest.action === "drop") {
				store.setPinned(varId, 0);
				store.setDropped(varId, 1, 0);
			} else if (latest.action === "unpin") {
				store.setPinned(varId, 0);
			} else {
				store.setPinned(varId, 0);
				store.setDropped(varId, 0, 1);
				if (latest.summary) store.updateSummary(varId, latest.summary.slice(0, cfg.summaryMaxChars));
			}
			for (const d of rows) consumedDecisionIds.add(d.id);
		}

		// ---- rollups ----
		let rollupCount = 0;
		if (initialRollupRows.length > 0 && parsed.rollup?.summary) {
			if (persistRollup(store, initialRollupRows, parsed.rollup.summary, cfg)) rollupCount++;
		}

		// If new archive decisions still push the store over the limit, roll up
		// the exact oldest excess during this same compaction.
		let activeAfter = store.listInContext().filter((v) => v.pinned === 0);
		if (activeAfter.length > cfg.archiveKeepIndividual) {
			const extraTargets = activeAfter.slice(0, activeAfter.length - cfg.archiveKeepIndividual);
			const generated = await callRollupAgent(ctx, cfg, extraTargets, signal);
			const fallback = extraTargets
				.map((v) => v.summary || cleanText(v.content, 200))
				.filter(Boolean)
				.join(" ");
			if (persistRollup(store, extraTargets, generated ?? fallback, cfg)) rollupCount++;
			activeAfter = store.listInContext().filter((v) => v.pinned === 0);
		}

		// ---- deterministic model-facing context ----
		const rendered = buildSummaryText(store, parsed.snapshot, cfg);
		// Decisions are consumed only once a valid final context exists.
		store.markDecisionsApplied([...consumedDecisionIds]);

		const details = {
			mode: "ctxvars-hybrid",
			candidates: candidates.length,
			decisions: parsed.decisions.length,
			pinned: store.listPinnedAll().length,
			archived: archivedCount,
			activeArchives: activeAfter.length,
			rollupsCreated: rollupCount,
			coreSummaryTokens: rendered.coreTokens,
			pinTokens: rendered.pinTokens,
			pinsOverWarningLimit: rendered.pinsOverWarningLimit,
			totalSummaryTokens: estimateTokens(rendered.text),
			storePath: store.path,
			ms: Date.now() - started,
		};
		dbg(`compaction: done in ${details.ms}ms ${JSON.stringify(details)}`);
		return { summary: rendered.text, usage: out.usage, details };
	} catch (err) {
		dbg(`compaction: error ${err instanceof Error ? err.message : String(err)}`);
		return null;
	}
}

function addDeps(store: VarStore, varId: number, refs: string[]) {
	let added = 0;
	for (const ref of refs) {
		const v = store.getVarByRef(ref) ?? (ref.startsWith("h_") || ref.startsWith("tool_") ? store.getByEntry(ref) : null);
		if (v) {
			store.addDep(varId, v.id);
			if (++added >= 5) break;
		}
	}
}

function clipAtWord(value: string, maxChars: number): string {
	if (value.length <= maxChars) return value;
	const raw = value.slice(0, Math.max(1, maxChars - 1));
	const boundary = raw.lastIndexOf(" ");
	return `${raw.slice(0, boundary > maxChars * 0.7 ? boundary : raw.length).trimEnd()}…`;
}

function cleanText(value: string, maxChars: number): string {
	return clipAtWord(value.replace(/\s+/g, " ").trim(), maxChars);
}

function validRefs(store: VarStore, refs: string[]): string[] {
	const seen = new Set<string>();
	return refs.filter((ref) => {
		if (seen.has(ref) || !store.getVarByRef(ref)) return false;
		seen.add(ref);
		return true;
	});
}

function formatItems(store: VarStore, items: SnapshotItem[], limit = 10): string[] {
	const seen = new Set<string>();
	const lines: string[] = [];
	for (const item of items) {
		const text = cleanText(item.text, 800);
		if (!text || seen.has(text.toLowerCase())) continue;
		seen.add(text.toLowerCase());
		const refs = validRefs(store, item.refs);
		lines.push(`- ${text}${refs.length ? ` [refs: ${refs.join(", ")}]` : ""}`);
		if (lines.length >= limit) break;
	}
	return lines;
}

function compactRanges(ids: number[]): string {
	if (ids.length === 0) return "";
	const sorted = [...new Set(ids)].sort((a, b) => a - b);
	const ranges: string[] = [];
	let start = sorted[0];
	let end = start;
	for (const id of sorted.slice(1)) {
		if (id === end + 1) {
			end = id;
			continue;
		}
		ranges.push(start === end ? `var_${start}` : `var_${start}..var_${end}`);
		start = end = id;
	}
	ranges.push(start === end ? `var_${start}` : `var_${start}..var_${end}`);
	return ranges.join(", ");
}

function renderCore(
	store: VarStore,
	snapshot: ContextSnapshot,
	topics: SnapshotTopic[],
	rollups: Array<{ range: string; summary: string; id: number }>,
): string {
	const lines: string[] = ["<context_snapshot>"];
	let task = cleanText(snapshot.currentTask, 1200);
	if (!task) {
		const latestUser = store.listRecent(200).find((v) => v.kind === "user" && v.dropped === 0);
		task = latestUser ? cleanText(latestUser.summary || latestUser.content, 1200) : "Continue the current task using the state below.";
	}
	lines.push("", "## Current task", task);

	const state = formatItems(store, snapshot.currentState, 12);
	if (state.length) lines.push("", "## Current state", ...state);

	const constraints = formatItems(store, snapshot.constraints, 10);
	const decisions = formatItems(store, snapshot.keyDecisions, 10);
	if (constraints.length || decisions.length) {
		lines.push("", "## Decisions and constraints");
		if (constraints.length) lines.push("", "### Constraints", ...constraints);
		if (decisions.length) lines.push("", "### Decisions", ...decisions);
	}

	const open = formatItems(store, snapshot.openItems, 10);
	if (open.length) lines.push("", "## Open work", ...open);

	if (topics.length) {
		lines.push("", "## Active topics");
		for (const topic of topics) {
			const name = cleanText(topic.name, 120);
			const summary = cleanText(topic.summary, 1800);
			if (!name || !summary) continue;
			const refs = validRefs(store, topic.refs);
			lines.push("", `### ${name}`, summary);
			if (refs.length) lines.push(`Refs: ${refs.join(", ")}`);
		}
	}

	const archives = store.listInContext().filter((v) => v.pinned === 0);
	if (rollups.length || archives.length) {
		lines.push("", "## Memory index");
		for (const r of rollups) lines.push(`- [${r.range}] ${cleanText(r.summary, 1200)}`);
		if (archives.length) {
			lines.push(`- ${archives.length} individually archived variable(s): ${compactRanges(archives.map((v) => v.id))}`);
		}
	}

	lines.push(
		"",
		"---",
		"Earlier content remains in the variable store. Use context_query to find it and context_read to load full content.",
		"</context_snapshot>",
	);
	return lines.join("\n");
}

/** Render a deterministic hybrid snapshot. The normal snapshot has a hard
 * budget; pins are appended separately and remain verbatim. */
function buildSummaryText(
	store: VarStore,
	snapshot: ContextSnapshot,
	cfg: CtxVarsConfig,
): { text: string; coreTokens: number; pinTokens: number; pinsOverWarningLimit: number } {
	let topics = snapshot.topics.slice(0, cfg.snapshotMaxTopics);
	let rollups = store.listRollups();
	// Defensive dedupe for stores created by older versions.
	const seenRanges = new Set<string>();
	rollups = rollups.filter((r) => {
		if (seenRanges.has(r.range)) return false;
		seenRanges.add(r.range);
		return true;
	});

	const maxChars = Math.max(1000, cfg.summaryTokenBudget * 4);
	let core = renderCore(store, snapshot, topics, rollups);
	// Remove low-priority detail first. Topics are ordered most important first;
	// rollups are ordered oldest first, so keep recent rollups when reducing.
	while (core.length > maxChars && topics.length > 3) {
		topics = topics.slice(0, -1);
		core = renderCore(store, snapshot, topics, rollups);
	}
	while (core.length > maxChars && rollups.length > 6) {
		rollups = rollups.slice(1);
		core = renderCore(store, snapshot, topics, rollups);
	}
	while (core.length > maxChars && topics.length > 0) {
		topics = topics.slice(0, -1);
		core = renderCore(store, snapshot, topics, rollups);
	}
	if (core.length > maxChars) {
		const suffix = "\n\n---\n[Snapshot reduced to configured token budget. Query the variable store for omitted details.]\n</context_snapshot>";
		core = core.slice(0, Math.max(0, maxChars - suffix.length)).trimEnd() + suffix;
	}

	const pins = store.listPinnedAll();
	let pinText = "";
	if (pins.length) {
		const lines = ["", "<pinned_context>", "## Pinned verbatim"];
		for (const p of pins) {
			lines.push("", `### var_${p.id} (${p.kind}, ${p.pinned === 2 ? "persistent" : "decayable"})`, p.content);
		}
		lines.push("</pinned_context>");
		pinText = lines.join("\n");
	}
	return {
		text: core + pinText,
		coreTokens: estimateTokens(core),
		pinTokens: estimateTokens(pinText),
		pinsOverWarningLimit: pins.filter((p) => p.content.length > cfg.pinnedMaxChars).length,
	};
}