import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";

const MESSAGE_TYPE = "context-sentinel";
const STATE_TYPE = "context-sentinel-state";

interface SentinelDetails {
	kind: "threshold";
	tier?: number;
	percent?: number;
	remainingTokens?: number;
}

interface SentinelState {
	highestAnnouncedTier: number;
}

const THRESHOLDS = [80, 90] as const;

interface UsageSnapshot {
	percent: number;
	remainingTokens?: number;
}

function getUsage(ctx: ExtensionContext): UsageSnapshot | undefined {
	const usage = ctx.getContextUsage();
	if (!usage || usage.percent === null || usage.percent === undefined) return undefined;
	return {
		percent: usage.percent,
		remainingTokens:
			usage.tokens === null ? undefined : Math.max(0, usage.contextWindow - usage.tokens),
	};
}

function formatTokens(count: number): string {
	return count >= 1000 ? `${Math.round(count / 1000)}k` : `${count}`;
}

function isSentinelDetails(value: unknown): value is SentinelDetails {
	if (!value || typeof value !== "object") return false;
	const details = value as Partial<SentinelDetails>;
	return details.kind === "threshold";
}

function isSentinelState(value: unknown): value is SentinelState {
	if (!value || typeof value !== "object") return false;
	return typeof (value as Partial<SentinelState>).highestAnnouncedTier === "number";
}

export default function contextSentinel(pi: ExtensionAPI) {
	let highestAnnouncedTier = 0;

	const announceCrossedTier = (ctx: ExtensionContext, deliverAs?: "steer"): void => {
		const usage = getUsage(ctx);
		if (!usage) return;

		const threshold = [...THRESHOLDS]
			.reverse()
			.find((candidate) => usage.percent >= candidate && candidate > highestAnnouncedTier);
		if (!threshold) return;

		highestAnnouncedTier = threshold;
		const remaining =
			usage.remainingTokens === undefined
				? ""
				: ` Approximately ${formatTokens(usage.remainingTokens)} tokens remain.`;
		pi.sendMessage(
			{
				customType: MESSAGE_TYPE,
				content: `[Context sentinel] Context window usage crossed ${threshold}%.${remaining}`,
				display: true,
				details: {
					kind: "threshold",
					tier: threshold,
					percent: usage.percent,
					remainingTokens: usage.remainingTokens,
				} satisfies SentinelDetails,
			},
			deliverAs ? { deliverAs } : undefined,
		);
		pi.appendEntry(STATE_TYPE, { highestAnnouncedTier } satisfies SentinelState);
	};

	pi.on("session_start", (_event, ctx) => {
		highestAnnouncedTier = 0;

		// Restore state from the active branch only. A compaction starts a fresh
		// pressure cycle because context usage drops after the summary is created.
		// State lives in custom entries (appendEntry); sentinel messages from
		// sessions created before the state entry existed still count as fallback.
		const branch = ctx.sessionManager.getBranch();
		for (let index = branch.length - 1; index >= 0; index--) {
			const entry = branch[index];
			if (entry.type === "compaction") break;
			if (entry.type === "custom" && entry.customType === STATE_TYPE && isSentinelState(entry.data)) {
				highestAnnouncedTier = Math.max(highestAnnouncedTier, entry.data.highestAnnouncedTier);
				continue;
			}
			if (entry.type !== "message" || entry.message.role !== "custom") continue;
			if (entry.message.customType !== MESSAGE_TYPE || !isSentinelDetails(entry.message.details)) continue;

			highestAnnouncedTier = Math.max(highestAnnouncedTier, entry.message.details.tier ?? 0);
		}

		// On resume/reload, surface any threshold already crossed before the user
		// submits another message.
		announceCrossedTier(ctx);
	});

	pi.on("turn_end", (event, ctx) => {
		if (event.toolResults.length === 0) return;

		// A tool-driven run already has another model call coming. Steer that
		// existing continuation as soon as context pressure crosses a tier.
		announceCrossedTier(ctx, "steer");
	});

	pi.on("agent_settled", (_event, ctx) => {
		// If a final response crossed the tier without another tool continuation,
		// append the warning while idle. It is visible now and reaches the model
		// with the user's next message, without manufacturing an extra turn.
		announceCrossedTier(ctx);
	});

	pi.on("session_compact", () => {
		// Pi owns compaction behavior. Sentinel only starts a fresh threshold cycle.
		highestAnnouncedTier = 0;
	});
}