import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { AssistantMessage, CacheRetention } from "@earendil-works/pi-ai";

export interface SessionCacheState {
	provider: string;
	modelId: string;
	cachedTokens: number;
	timestamp: number;
	ttlSeconds?: number;
	deterministic: boolean;
}

export interface PromptCacheSnapshot {
	cachedTokens: number;
	ttlSeconds?: number;
	expiresAt?: number;
	deterministic: boolean;
	status: "warm" | "cold" | "mismatch";
	canRestore?: boolean;
}

export type CacheTtlStatus =
	| { status: "none" }
	| { status: "cold"; cachedTokens?: number }
	| { status: "mismatch"; cachedModelId: string; canRestore: boolean }
	| {
			status: "warm";
			cachedTokens: number;
			ttlSeconds: number;
			expiresAt: number;
			remainingSeconds: number;
			deterministic: boolean;
			label: string;
	  };

export interface ModelIdentity {
	provider: string;
	id: string;
}

export function formatCacheTtl(seconds: number): string {
	if (!Number.isFinite(seconds) || seconds <= 0) return "0s";
	const total = Math.floor(seconds);
	if (total <= 0) return "0s";
	if (total < 60) return `${total}s`;
	const minutes = Math.floor(total / 60);
	const remSeconds = total % 60;
	if (total < 3600) {
		return remSeconds > 0 ? `${minutes}m ${remSeconds}s` : `${minutes}m`;
	}
	const hours = Math.floor(total / 3600);
	const remMinutes = Math.floor((total % 3600) / 60);
	return remMinutes > 0 ? `${hours}h ${remMinutes}m` : `${hours}h`;
}

export function resolvePromptCacheTtl(
	provider: string,
	modelId: string,
	cacheRetention?: CacheRetention,
	supportsLongCacheRetention = true,
): { ttlSeconds: number; deterministic: boolean } | undefined {
	const prov = provider.toLowerCase();
	const id = modelId.toLowerCase();
	const isAnthropic = prov === "anthropic" || id.startsWith("claude-");
	const isBedrockClaude = prov === "amazon-bedrock" && id.includes("claude");
	if (isAnthropic || isBedrockClaude) {
		const isLong =
			supportsLongCacheRetention &&
			(cacheRetention === "long" || (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long"));
		return { ttlSeconds: isLong ? 3600 : 300, deterministic: true };
	}

	if (prov === "openai-codex") {
		return { ttlSeconds: 1800, deterministic: true };
	}

	const isOpenAI = prov === "openai" || id.startsWith("gpt-") || prov.includes("azure-openai");
	if (isOpenAI) {
		const isLong =
			supportsLongCacheRetention &&
			(cacheRetention === "long" || (typeof process !== "undefined" && process.env.PI_CACHE_RETENTION === "long"));
		if (isLong) return { ttlSeconds: 86400, deterministic: true };
		return { ttlSeconds: 1800, deterministic: true };
	}

	if (prov === "deepseek" || id.startsWith("deepseek-")) {
		return { ttlSeconds: 43200, deterministic: false };
	}

	if (prov.startsWith("google") || id.startsWith("gemini-")) {
		return { ttlSeconds: 300, deterministic: false };
	}

	if (prov === "groq") {
		return { ttlSeconds: 7200, deterministic: false };
	}

	if (prov === "zai" || prov === "zhipu" || id.startsWith("glm-")) {
		return { ttlSeconds: 1800, deterministic: false };
	}

	if (prov === "mistral" || id.startsWith("mistral-") || id.startsWith("codestral-")) {
		return { ttlSeconds: 1800, deterministic: false };
	}

	if (prov === "together" || prov === "together-ai") {
		return { ttlSeconds: 1800, deterministic: false };
	}

	if (prov === "fireworks") {
		return { ttlSeconds: 900, deterministic: false };
	}

	return undefined;
}

export function findLatestCompletedAssistant(
	messages: readonly AgentMessage[],
	latestCompactionTimestamp?: number,
): AssistantMessage | undefined {
	for (let i = messages.length - 1; i >= 0; i--) {
		const msg = messages[i];
		if (msg.role !== "assistant" || !("usage" in msg)) continue;
		const assistant = msg as AssistantMessage;
		if (assistant.stopReason === "aborted" || assistant.stopReason === "error" || !assistant.usage) continue;
		if (latestCompactionTimestamp !== undefined && assistant.timestamp <= latestCompactionTimestamp) {
			return undefined;
		}
		return assistant;
	}
	return undefined;
}

export function extractSessionCacheState(
	messages: readonly AgentMessage[],
	latestCompactionTimestamp?: number,
	cacheRetention?: CacheRetention,
	supportsLongCacheRetention?: boolean,
): SessionCacheState | undefined {
	const assistant = findLatestCompletedAssistant(messages, latestCompactionTimestamp);
	if (!assistant || !assistant.usage) return undefined;
	const cachedTokens = (assistant.usage.cacheRead ?? 0) + (assistant.usage.cacheWrite ?? 0);
	if (cachedTokens <= 0) return undefined;

	const provider = assistant.provider || "unknown";
	const modelId = assistant.model || "unknown";
	const ttlInfo = resolvePromptCacheTtl(provider, modelId, cacheRetention, supportsLongCacheRetention);

	return {
		provider,
		modelId,
		cachedTokens,
		timestamp: assistant.timestamp,
		ttlSeconds: ttlInfo?.ttlSeconds,
		deterministic: ttlInfo?.deterministic ?? false,
	};
}

export function evaluateActiveCacheTtl(
	currentModel: ModelIdentity | undefined,
	cache: SessionCacheState | undefined,
	now: number = Date.now(),
): CacheTtlStatus {
	if (!currentModel || !cache || cache.cachedTokens <= 0) {
		return { status: "none" };
	}

	if (
		currentModel.provider.toLowerCase() !== cache.provider.toLowerCase() ||
		currentModel.id.toLowerCase() !== cache.modelId.toLowerCase()
	) {
		const canRestore = cache.ttlSeconds !== undefined && cache.timestamp + cache.ttlSeconds * 1000 > now;
		return {
			status: "mismatch",
			cachedModelId: cache.modelId,
			canRestore,
		};
	}

	if (cache.ttlSeconds === undefined) {
		return {
			status: "warm",
			cachedTokens: cache.cachedTokens,
			ttlSeconds: 0,
			expiresAt: 0,
			remainingSeconds: 0,
			deterministic: false,
			label: "",
		};
	}

	const expiresAt = cache.timestamp + cache.ttlSeconds * 1000;
	const remainingSeconds = Math.max(0, Math.floor((expiresAt - now) / 1000));
	if (remainingSeconds <= 0) {
		return { status: "cold", cachedTokens: cache.cachedTokens };
	}

	return {
		status: "warm",
		cachedTokens: cache.cachedTokens,
		ttlSeconds: cache.ttlSeconds,
		expiresAt,
		remainingSeconds,
		deterministic: cache.deterministic,
		label: formatCacheTtl(remainingSeconds),
	};
}

export {
	assertPromptPayloadWithinCeiling,
	calculatePromptPayloadCeiling,
	compilePromptPayload,
	DEFAULT_PROMPT_PAYLOAD_CEILING_RATIO,
	estimateMessageTokens,
	estimatePayloadTokens,
	formatModelContextCeilingStatement,
	PromptPayloadCeilingExceededError,
} from "./prompts/payload-ceiling.js";