import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, CacheRetention, ProviderHeaders, Tool, Transport, Usage } from "@earendil-works/pi-ai";
import { convertToLlm, type ExtensionAPI, type ExtensionContext } from "@earendil-works/pi-coding-agent";

const EXTENSION_NAME = "Cachepoint";
const OPENAI_PROVIDER = "openai";
const OPENAI_CODEX_PROVIDER = "openai-codex";
const ANTHROPIC_PROVIDER = "anthropic";
const KIMI_CODING_PROVIDER = "kimi-coding";
const SUPPORTED_PROVIDERS = new Set([
	OPENAI_PROVIDER,
	OPENAI_CODEX_PROVIDER,
	ANTHROPIC_PROVIDER,
	KIMI_CODING_PROVIDER,
]);
const STATUS_KEY = "cachepoint";
const STATUS_COUNTDOWN_WINDOW_MS = 60 * 1000;
const STATUS_UPDATE_INTERVAL_MS = 1000;
const SHORT_CACHE_TTL_MS = 5 * 60 * 1000;
const OPENAI_LONG_CACHE_TTL_MS = 24 * 60 * 60 * 1000;
const ANTHROPIC_LONG_CACHE_TTL_MS = 60 * 60 * 1000;
const SHORT_CACHE_MARGIN_MS = 60 * 1000;
const LONG_CACHE_MARGIN_MS = 5 * 60 * 1000;

const SUMMARY_INSTRUCTIONS = `This is an automatic session checkpoint request.

Do not continue the task, make changes, or call any tools. Return only a self-contained continuation summary of the current session as structured Markdown.

The summary will replace older conversation history, while a recent tail of the conversation will also be retained verbatim. Preserve everything needed to resume accurately, including:

## Goal
- The user's objective and intended outcome

## Constraints & Preferences
- Explicit requirements, preferences, and things that must not be changed

## Progress
### Done
- Completed work and verified results

### In Progress
- Current work and its exact state

### Blocked
- Errors, unresolved issues, and missing information

## Key Decisions
- Decisions made and their rationale

## Important Technical Context
- Relevant architecture, APIs, commands, identifiers, values, and discoveries
- Files read and files modified, with the important changes

## Next Steps
1. Concrete actions required to continue

Be concise, but favor correctness and continuity over brevity. Do not mention this checkpoint request.`;

interface Config {
	minTokens: number;
	maxSummaryTokens: number;
	debug: boolean;
}

interface RequestSnapshot {
	modelKey: string;
	systemPrompt: string;
	messages: AgentMessage[];
	tools: Tool[];
	providerHeaders?: ProviderHeaders;
	providerPayload?: unknown;
	transport?: Transport;
}

interface CachepointDetails {
	extension: "pi-cachepoint";
	provider: string;
	model: string;
	generatedAt: string;
	cacheRetention: CacheRetention;
	cacheRead: number;
	cacheWrite: number;
	input: number;
	output: number;
}

function numberFlag(pi: ExtensionAPI, name: string, fallback: number, minimum = 0): number {
	const raw = pi.getFlag(name);
	const parsed = typeof raw === "string" ? Number(raw) : Number.NaN;
	return Number.isFinite(parsed) && parsed >= minimum ? parsed : fallback;
}

function getConfig(pi: ExtensionAPI): Config {
	return {
		minTokens: Math.floor(numberFlag(pi, "cachepoint-min-tokens", 50_000, 1)),
		maxSummaryTokens: Math.floor(numberFlag(pi, "cachepoint-max-summary-tokens", 8192, 256)),
		debug: pi.getFlag("cachepoint-debug") === true,
	};
}

function modelKey(ctx: ExtensionContext): string | undefined {
	return ctx.model ? `${ctx.model.provider}/${ctx.model.id}` : undefined;
}

function isSupportedProvider(provider: string | undefined): boolean {
	return provider !== undefined && SUPPORTED_PROVIDERS.has(provider);
}

function isSupported(ctx: ExtensionContext): boolean {
	return isSupportedProvider(ctx.model?.provider);
}

function clone<T>(value: T): T {
	return structuredClone(value);
}

function snapshotTools(pi: ExtensionAPI): Tool[] {
	const byName = new Map(pi.getAllTools().map((tool) => [tool.name, tool]));
	return pi.getActiveTools().map((name) => {
		const tool = byName.get(name);
		if (!tool) throw new Error(`Active tool definition not found: ${name}`);
		return {
			name: tool.name,
			description: tool.description,
			parameters: tool.parameters,
		};
	});
}

function textFrom(response: AssistantMessage): string {
	return response.content
		.filter((block): block is { type: "text"; text: string } => block.type === "text")
		.map((block) => block.text)
		.join("\n")
		.trim();
}

function promptTokens(usage: Usage): number {
	return usage.input + usage.cacheRead + usage.cacheWrite;
}

function formatDuration(milliseconds: number): string {
	const seconds = Math.max(0, Math.ceil(milliseconds / 1000));
	if (seconds < 60) return `${seconds}s`;
	const minutes = Math.floor(seconds / 60);
	const remainingSeconds = seconds % 60;
	if (minutes < 60) return remainingSeconds ? `${minutes}m ${remainingSeconds}s` : `${minutes}m`;
	const hours = Math.floor(minutes / 60);
	const remainingMinutes = minutes % 60;
	return remainingMinutes ? `${hours}h ${remainingMinutes}m` : `${hours}h`;
}

function cacheRetentionFromEnv(env?: Record<string, string>): CacheRetention {
	return (env?.PI_CACHE_RETENTION ?? process.env.PI_CACHE_RETENTION) === "long" ? "long" : "short";
}

function explicitlyDisablesLongRetention(model: { compat?: unknown }): boolean {
	return (
		typeof model.compat === "object" &&
		model.compat !== null &&
		(model.compat as Record<string, unknown>).supportsLongCacheRetention === false
	);
}

function effectiveCacheRetention(
	model: { provider: string; compat?: unknown },
	env?: Record<string, string>,
): CacheRetention {
	const requested = cacheRetentionFromEnv(env);
	// Codex and Kimi Coding can reuse prompt prefixes, but neither transport
	// exposes a caller-controlled long-retention policy.
	if (
		requested === "long" &&
		(model.provider === OPENAI_CODEX_PROVIDER ||
			model.provider === KIMI_CODING_PROVIDER ||
			explicitlyDisablesLongRetention(model))
	) {
		return "short";
	}
	return requested;
}

function cacheTiming(
	provider: string | undefined,
	retention: CacheRetention,
): { ttlMs: number; marginMs: number } {
	const ttlMs =
		retention === "long"
			? provider === ANTHROPIC_PROVIDER
				? ANTHROPIC_LONG_CACHE_TTL_MS
				: OPENAI_LONG_CACHE_TTL_MS
			: SHORT_CACHE_TTL_MS;
	const requestedMargin = retention === "long" ? LONG_CACHE_MARGIN_MS : SHORT_CACHE_MARGIN_MS;
	return { ttlMs, marginMs: Math.min(requestedMargin, ttlMs - 1000) };
}

export default function cachepoint(pi: ExtensionAPI) {
	pi.registerFlag("cachepoint-min-tokens", {
		description: "Minimum current context tokens before automatic Cachepoint compaction",
		type: "string",
		default: "50000",
	});
	pi.registerFlag("cachepoint-max-summary-tokens", {
		description: "Maximum output-token budget for a Cachepoint summary",
		type: "string",
		default: "8192",
	});
	pi.registerFlag("cachepoint-debug", {
		description: "Show Cachepoint scheduling diagnostics",
		type: "boolean",
		default: false,
	});

	let timer: ReturnType<typeof setTimeout> | undefined;
	let statusTimer: ReturnType<typeof setTimeout> | undefined;
	let timerVersion = 0;
	let closed = false;
	let cacheActivityAt: number | undefined;
	let requestStartedAt: number | undefined;
	let retention: CacheRetention = "short";
	let pendingProviderHeaders: ProviderHeaders | undefined;
	let requestSnapshot: RequestSnapshot | undefined;
	let latestAssistant: AssistantMessage | undefined;
	let cachepointPending = false;
	let cachepointRunning = false;

	const clearTimer = () => {
		if (timer) clearTimeout(timer);
		if (statusTimer) clearTimeout(statusTimer);
		timer = undefined;
		statusTimer = undefined;
		timerVersion++;
	};

	const clearSchedule = (ctx?: ExtensionContext) => {
		clearTimer();
		cacheActivityAt = undefined;
		requestStartedAt = undefined;
		pendingProviderHeaders = undefined;
		if (ctx?.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
	};

	const debug = (ctx: ExtensionContext, message: string) => {
		if (getConfig(pi).debug && ctx.hasUI) ctx.ui.notify(`${EXTENSION_NAME}: ${message}`, "info");
	};

	const resolveRetention = async (ctx: ExtensionContext): Promise<CacheRetention> => {
		if (!ctx.model) return "short";
		const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
		return effectiveCacheRetention(ctx.model, auth.ok ? auth.env : undefined);
	};

	const currentDeadline = (ctx: ExtensionContext): number | undefined => {
		if (cacheActivityAt === undefined) return undefined;
		const timing = cacheTiming(ctx.model?.provider, retention);
		return cacheActivityAt + timing.ttlMs - timing.marginMs;
	};

	const currentExpiry = (ctx: ExtensionContext): number | undefined => {
		if (cacheActivityAt === undefined) return undefined;
		return cacheActivityAt + cacheTiming(ctx.model?.provider, retention).ttlMs;
	};

	const refreshCountdownStatus = (ctx: ExtensionContext, version: number) => {
		if (closed || version !== timerVersion || !ctx.hasUI) return;
		statusTimer = undefined;

		const config = getConfig(pi);
		const deadline = currentDeadline(ctx);
		if (!isSupported(ctx) || cachepointRunning || deadline === undefined) {
			ctx.ui.setStatus(STATUS_KEY, undefined);
			return;
		}

		const remaining = deadline - Date.now();
		const tokens = ctx.getContextUsage()?.tokens;
		const thresholdReached =
			tokens !== null && tokens !== undefined && tokens >= config.minTokens;
		if (remaining < STATUS_COUNTDOWN_WINDOW_MS && thresholdReached) {
			ctx.ui.setStatus(STATUS_KEY, `cachepoint in ${formatDuration(remaining)}`);
		} else {
			ctx.ui.setStatus(STATUS_KEY, undefined);
		}

		if (remaining <= 0) return;
		const nextUpdateMs =
			remaining >= STATUS_COUNTDOWN_WINDOW_MS
				? remaining - STATUS_COUNTDOWN_WINDOW_MS + 1
				: Math.min(STATUS_UPDATE_INTERVAL_MS, remaining);
		statusTimer = setTimeout(() => refreshCountdownStatus(ctx, version), nextUpdateMs);
		statusTimer.unref?.();
	};

	const runCachepoint = (ctx: ExtensionContext) => {
		if (cachepointRunning) return;
		cachepointRunning = true;
		cachepointPending = true;
		clearTimer();

		if (ctx.hasUI) {
			ctx.ui.setStatus(STATUS_KEY, "cachepoint: summarizing");
			ctx.ui.notify(`${EXTENSION_NAME}: creating cache-aware summary`, "info");
		}

		ctx.compact({
			onComplete: (result) => {
				cachepointRunning = false;
				cachepointPending = false;
				cacheActivityAt = undefined;
				if (ctx.hasUI) {
					ctx.ui.setStatus(STATUS_KEY, undefined);
					const usage = result.usage;
					const estimatedTokens = result.estimatedTokensAfter?.toLocaleString() ?? "unknown";
					const cacheMessage = usage
						? `${usage.cacheRead.toLocaleString()} cache-read tokens; ${estimatedTokens} estimated tokens remain`
						: `${estimatedTokens} estimated tokens remain`;
					ctx.ui.notify(`${EXTENSION_NAME} created: ${cacheMessage}`, "info");
				}
			},
			onError: (error) => {
				cachepointRunning = false;
				cachepointPending = false;
				cacheActivityAt = undefined;
				if (ctx.hasUI) {
					ctx.ui.setStatus(STATUS_KEY, undefined);
					ctx.ui.notify(`${EXTENSION_NAME} was not created: ${error.message}`, "warning");
				}
			},
		});
	};

	const handleDeadline = async (ctx: ExtensionContext, version: number) => {
		if (closed || version !== timerVersion) return;
		timer = undefined;

		const config = getConfig(pi);
		if (!isSupported(ctx) || cacheActivityAt === undefined) {
			clearSchedule(ctx);
			return;
		}

		const expiry = currentExpiry(ctx);
		if (expiry === undefined || Date.now() >= expiry) {
			debug(ctx, "assumed cache TTL expired before the checkpoint could run");
			clearSchedule(ctx);
			return;
		}

		if (!ctx.isIdle() || ctx.hasPendingMessages()) {
			debug(ctx, "deadline reached while busy; waiting for agent_settled");
			return;
		}

		const usage = ctx.getContextUsage();
		if (!usage || usage.tokens === null || usage.tokens < config.minTokens) {
			debug(ctx, `deadline reached, but context is below ${config.minTokens.toLocaleString()} tokens`);
			clearSchedule(ctx);
			return;
		}

		if (!requestSnapshot || requestSnapshot.modelKey !== modelKey(ctx) || latestAssistant?.stopReason !== "stop") {
			debug(ctx, "deadline reached without a reusable completed request snapshot");
			clearSchedule(ctx);
			return;
		}

		const summaryBudget = Math.min(config.maxSummaryTokens, ctx.model?.maxTokens || config.maxSummaryTokens);
		if (usage.tokens + summaryBudget + 1024 >= usage.contextWindow) {
			if (ctx.hasUI) {
				ctx.ui.notify(`${EXTENSION_NAME}: insufficient context headroom for a cache-aware summary`, "warning");
			}
			clearSchedule(ctx);
			return;
		}

		runCachepoint(ctx);
	};

	const armTimer = async (ctx: ExtensionContext) => {
		clearTimer();
		if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
		const version = timerVersion;
		if (closed || !isSupported(ctx) || cacheActivityAt === undefined || cachepointRunning) {
			if (ctx.hasUI) ctx.ui.setStatus(STATUS_KEY, undefined);
			return;
		}

		try {
			retention = await resolveRetention(ctx);
		} catch {
			retention = ctx.model ? effectiveCacheRetention(ctx.model) : cacheRetentionFromEnv();
		}
		if (closed || version !== timerVersion) return;

		const expiry = currentExpiry(ctx);
		if (expiry === undefined || Date.now() >= expiry) {
			debug(ctx, "assumed cache TTL expired while waiting for pi to become idle");
			clearSchedule(ctx);
			return;
		}

		const deadline = currentDeadline(ctx);
		if (deadline === undefined) return;
		const delay = Math.max(0, deadline - Date.now());
		debug(ctx, `armed for ${formatDuration(delay)} (${retention} retention)`);

		timer = setTimeout(() => {
			void handleDeadline(ctx, version);
		}, delay);
		timer.unref?.();
		refreshCountdownStatus(ctx, version);
	};

	pi.on("session_start", (_event, ctx) => {
		closed = false;
		cachepointPending = false;
		cachepointRunning = false;
		requestSnapshot = undefined;
		latestAssistant = undefined;
		clearSchedule(ctx);
	});

	pi.on("context", (event, ctx) => {
		if (!isSupported(ctx)) return;
		try {
			requestSnapshot = {
				modelKey: modelKey(ctx)!,
				systemPrompt: ctx.getSystemPrompt(),
				messages: clone(event.messages),
				tools: snapshotTools(pi),
			};
		} catch (error) {
			requestSnapshot = undefined;
			debug(ctx, `could not snapshot request context: ${error instanceof Error ? error.message : String(error)}`);
		}
	});

	pi.on("input", (event, ctx) => {
		// Cancel at the earliest user-input lifecycle point so a deadline that
		// elapsed during the new request cannot compact when the agent settles.
		// Extension-injected messages are not user activity and must not disturb
		// the schedule (in particular, Cachepoint's own internal work).
		if (event.source === "extension" || !isSupported(ctx) || cachepointRunning) return;
		if (cacheActivityAt === undefined) return;
		debug(ctx, "user input received; waiting for the next successful provider response");
		clearSchedule(ctx);
	});

	pi.on("before_provider_headers", (event, ctx) => {
		if (!isSupported(ctx) || cachepointRunning) return;
		// Keep the live object until before_provider_request so mutations made by
		// later header handlers are included in the snapshot.
		pendingProviderHeaders = event.headers;
	});

	pi.on("before_provider_request", (event, ctx) => {
		if (!isSupported(ctx) || cachepointRunning) return;
		requestStartedAt = Date.now();
		const snapshot = requestSnapshot;
		if (snapshot && snapshot.modelKey === modelKey(ctx)) {
			try {
				snapshot.providerHeaders = pendingProviderHeaders ? clone(pendingProviderHeaders) : undefined;
				snapshot.providerPayload = clone(event.payload);
			} catch {
				snapshot.providerHeaders = undefined;
				snapshot.providerPayload = undefined;
			}
		}
		pendingProviderHeaders = undefined;
	});

	pi.on("after_provider_response", (event, ctx) => {
		if (!isSupported(ctx) || cachepointRunning || event.status < 200 || event.status >= 300) return;
		cacheActivityAt = requestStartedAt ?? Date.now();
		requestStartedAt = undefined;
		const snapshot = requestSnapshot;
		if (ctx.model?.provider === OPENAI_CODEX_PROVIDER && snapshot && snapshot.modelKey === modelKey(ctx)) {
			snapshot.transport = "sse";
		}
		void armTimer(ctx);
	});

	pi.on("message_end", (event, ctx) => {
		if (event.message.role !== "assistant" || !isSupportedProvider(event.message.provider)) return;
		latestAssistant = clone(event.message);

		// Successful Codex WebSocket requests do not have an HTTP response and
		// therefore do not emit after_provider_response. Treat the finalized
		// assistant message as the successful response signal in that case. SSE
		// Codex requests already cleared requestStartedAt in the response hook.
		if (event.message.provider === OPENAI_CODEX_PROVIDER && requestStartedAt !== undefined) {
			const startedAt = requestStartedAt;
			requestStartedAt = undefined;
			if (
				event.message.stopReason !== "pending" &&
				event.message.stopReason !== "error" &&
				event.message.stopReason !== "aborted"
			) {
				cacheActivityAt = startedAt;
				void armTimer(ctx);
			}
		}
	});

	pi.on("agent_settled", (_event, ctx) => {
		void armTimer(ctx);
	});

	pi.on("model_select", (_event, ctx) => {
		requestSnapshot = undefined;
		latestAssistant = undefined;
		clearSchedule(ctx);
	});

	pi.on("session_compact", (_event, ctx) => {
		requestSnapshot = undefined;
		latestAssistant = undefined;
		clearSchedule(ctx);
	});

	pi.on("session_tree", (_event, ctx) => {
		requestSnapshot = undefined;
		latestAssistant = undefined;
		clearSchedule(ctx);
	});

	pi.on("session_shutdown", (_event, ctx) => {
		closed = true;
		clearSchedule(ctx);
	});

	pi.on("session_before_compact", async (event, ctx) => {
		if (!cachepointPending) return;
		cachepointPending = false;

		try {
			const model = ctx.model;
			const snapshot = requestSnapshot;
			const assistant = latestAssistant;
			if (!model || !isSupportedProvider(model.provider)) {
				throw new Error("the active model is not from a supported provider");
			}
			if (!snapshot || snapshot.modelKey !== modelKey(ctx)) {
				throw new Error("the current provider request snapshot is unavailable");
			}
			if (!assistant || assistant.stopReason !== "stop") {
				throw new Error("the latest provider turn did not complete normally");
			}

			const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
			if (auth.ok === false) throw new Error(auth.error);
			const provider = ctx.modelRegistry.getProvider(model.provider);
			if (!provider) throw new Error("the effective provider is unavailable");

			const effectiveRetention = effectiveCacheRetention(model, auth.env);
			const config = getConfig(pi);
			const maxTokens = Math.min(config.maxSummaryTokens, model.maxTokens || config.maxSummaryTokens);
			const previousPayload = snapshot.providerPayload;
			const context = {
				systemPrompt: snapshot.systemPrompt,
				messages: [
					...convertToLlm(snapshot.messages),
					assistant,
					{
						role: "user" as const,
						content: [{ type: "text" as const, text: SUMMARY_INSTRUCTIONS }],
						timestamp: Date.now(),
					},
				],
				tools: snapshot.tools.length ? snapshot.tools : undefined,
			};

			const response = await provider
				.streamSimple(model, context, {
					apiKey: auth.apiKey,
					headers: snapshot.providerHeaders ?? auth.headers,
					env: auth.env,
					signal: event.signal,
					transport: snapshot.transport,
					reasoning: ctx.thinkingLevel === "off" ? undefined : ctx.thinkingLevel,
					maxTokens,
					cacheRetention: effectiveRetention,
					sessionId: ctx.sessionManager.getSessionId(),
					onPayload: (payload) => {
						if (!payload || typeof payload !== "object") return payload;
						const next = { ...(payload as Record<string, unknown>) };
						if (previousPayload && typeof previousPayload === "object") {
							const previous = previousPayload as Record<string, unknown>;
							if ("tools" in previous) next.tools = clone(previous.tools);
						}
						if (Array.isArray(next.tools) && next.tools.length > 0) {
							next.tool_choice =
								model.api === "anthropic-messages" ? { type: "none" } : "none";
						}
						return next;
					},
				})
				.result();

			if (event.signal.aborted) throw new Error("summary request was cancelled");
			if (response.stopReason !== "stop") {
				throw new Error(response.errorMessage || `summary stopped with reason ${response.stopReason}`);
			}
			if (response.content.some((block) => block.type === "toolCall")) {
				throw new Error("the summary response attempted to call a tool");
			}

			const summary = textFrom(response);
			if (!summary) throw new Error("the summary response was empty");

			if (ctx.hasUI && response.usage.cacheRead === 0) {
				ctx.ui.notify(
					`${EXTENSION_NAME}: ${model.provider} reported no cache read for the summary request`,
					"warning",
				);
			}
			debug(
				ctx,
				`summary prompt ${promptTokens(response.usage).toLocaleString()} tokens; cache read ${response.usage.cacheRead.toLocaleString()}`,
			);

			const details: CachepointDetails = {
				extension: "pi-cachepoint",
				provider: model.provider,
				model: model.id,
				generatedAt: new Date().toISOString(),
				cacheRetention: effectiveRetention,
				cacheRead: response.usage.cacheRead,
				cacheWrite: response.usage.cacheWrite,
				input: response.usage.input,
				output: response.usage.output,
			};

			return {
				compaction: {
					summary,
					firstKeptEntryId: event.preparation.firstKeptEntryId,
					tokensBefore: event.preparation.tokensBefore,
					usage: response.usage,
					details,
				},
			};
		} catch (error) {
			const message = error instanceof Error ? error.message : String(error);
			if (ctx.hasUI && !event.signal.aborted) ctx.ui.notify(`${EXTENSION_NAME}: ${message}`, "error");
			return { cancel: true };
		}
	});

	pi.registerCommand("cachepoint-status", {
		description: "Show Cachepoint provider, context, and timer status",
		handler: async (_args, ctx) => {
			const config = getConfig(pi);
			const usage = ctx.getContextUsage();
			let effectiveRetention = retention;
			try {
				effectiveRetention = await resolveRetention(ctx);
			} catch {
				// Keep the last known/default value.
			}
			const deadline = currentDeadline(ctx);
			const lines = [
				`Provider: ${modelKey(ctx) ?? "none"}${isSupported(ctx) ? " (supported)" : " (unsupported)"}`,
				`Context: ${usage?.tokens === null || usage?.tokens === undefined ? "unknown" : usage.tokens.toLocaleString()} / ${usage?.contextWindow?.toLocaleString() ?? "unknown"} tokens`,
				`Minimum: ${config.minTokens.toLocaleString()} tokens`,
				`Retention policy: ${effectiveRetention}`,
				`Timer: ${deadline === undefined ? "not armed" : formatDuration(deadline - Date.now())}`,
				`State: ${cachepointRunning ? "summarizing" : "idle"}`,
			];
			ctx.ui.notify(lines.join("\n"), "info");
		},
	});

}