import { SPAN_KEYS } from "../constants/otel-attributes.js";
import { GATEWAY_ROUTABLE_PROVIDERS_SET } from "../constants/models.js";
import { getLogger } from "./logger.js";
import { createHash } from "node:crypto";
import { trace } from "@opentelemetry/api";

//#region src/utils/usage-cost-middleware.ts
const logger = getLogger("usage-cost-middleware");
function extractUsageTokens(usage) {
	const inputTokens = typeof usage?.inputTokens === "object" ? usage.inputTokens.total ?? 0 : usage?.inputTokens ?? 0;
	const outputTokens = typeof usage?.outputTokens === "object" ? usage.outputTokens.total ?? 0 : usage?.outputTokens ?? 0;
	const reasoningTokens = typeof usage?.outputTokens === "object" ? usage.outputTokens.reasoning : void 0;
	let cachedReadTokens = typeof usage?.inputTokens === "object" ? usage.inputTokens.cacheRead : void 0;
	const cachedWriteTokens = typeof usage?.inputTokens === "object" ? usage.inputTokens.cacheWrite : void 0;
	if (cachedReadTokens === void 0) {
		const geminiCacheRead = usage?.inputTokenDetails?.cacheReadTokens ?? usage?.cachedInputTokens;
		if (typeof geminiCacheRead === "number") cachedReadTokens = geminiCacheRead;
	}
	return {
		inputTokens,
		outputTokens,
		reasoningTokens,
		cachedReadTokens,
		cachedWriteTokens
	};
}
function extractGatewayCost(providerMetadata) {
	const gw = providerMetadata?.gateway;
	if (!gw) return 0;
	const cost = parseFloat(gw.cost);
	if (!Number.isNaN(cost) && cost > 0) return cost;
	const marketCost = parseFloat(gw.marketCost);
	if (!Number.isNaN(marketCost) && marketCost > 0) return marketCost;
	return 0;
}
function normalizeModelId(modelId) {
	const slashIndex = modelId.indexOf("/");
	if (slashIndex === -1) return modelId;
	const prefix = modelId.slice(0, slashIndex);
	if (GATEWAY_ROUTABLE_PROVIDERS_SET.has(prefix)) return modelId.slice(slashIndex + 1);
	return modelId;
}
function setGatewayAttributesOnSpan(providerMetadata) {
	const activeSpan = trace.getActiveSpan();
	if (!activeSpan) return;
	const cost = extractGatewayCost(providerMetadata);
	activeSpan.setAttribute(SPAN_KEYS.GEN_AI_COST_ESTIMATED_USD, cost);
	const gw = providerMetadata?.gateway;
	if (gw) {
		if (cost === 0) logger.warn({ gateway: gw }, "Routed through gateway but no cost data in response");
		const routing = gw.routing;
		if (routing?.finalProvider) activeSpan.setAttribute(SPAN_KEYS.GEN_AI_RESPONSE_PROVIDER, routing.finalProvider);
		if (routing?.resolvedProvider) activeSpan.setAttribute(SPAN_KEYS.GEN_AI_REQUEST_PROVIDER, routing.resolvedProvider);
	}
}
function countCacheMarkers(prompt, callProviderOptions) {
	let count = 0;
	const caching = callProviderOptions?.gateway?.caching;
	if (caching && caching !== "off" && caching !== "disabled") count++;
	for (const msg of prompt) {
		if (msg.providerOptions?.anthropic?.cacheControl) count++;
		if (Array.isArray(msg.content)) {
			for (const part of msg.content) if (part?.providerOptions?.anthropic?.cacheControl) count++;
		}
	}
	return Math.min(count, 4);
}
/**
* Count of part-level (conversation-history / BP2) cache markers only — the distinct BP2 signal.
* BP1 (system + tools) is message-level and tracked via the prefix signature; this isolates the
* history boundary so cross-turn cache participation can be attributed to it.
*/
function countHistoryCacheMarkers(prompt) {
	let count = 0;
	for (const msg of prompt) if (Array.isArray(msg.content)) {
		for (const part of msg.content) if (part?.providerOptions?.anthropic?.cacheControl) count++;
	}
	return count;
}
function computePrefixSignature(prompt, tools) {
	const systemParts = [];
	for (const msg of prompt) {
		if (msg.role !== "system") continue;
		if (typeof msg.content === "string") systemParts.push(msg.content);
		else if (Array.isArray(msg.content)) {
			for (const part of msg.content) if (typeof part === "string") systemParts.push(part);
			else if (part?.text) systemParts.push(part.text);
		}
	}
	const toolParts = (tools ?? []).map((tool) => [
		tool.name ?? "",
		tool.description ?? "",
		tool.parameters ? JSON.stringify(tool.parameters) : ""
	]);
	return createHash("sha256").update(JSON.stringify([systemParts, toolParts])).digest("hex").slice(0, 10);
}
function setCacheAttributesOnSpan(params, usage) {
	const activeSpan = trace.getActiveSpan();
	if (!activeSpan) return;
	const { cachedReadTokens, cachedWriteTokens } = extractUsageTokens(usage);
	activeSpan.setAttribute(SPAN_KEYS.GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS, typeof cachedReadTokens === "number" ? cachedReadTokens : 0);
	activeSpan.setAttribute(SPAN_KEYS.GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS, typeof cachedWriteTokens === "number" ? cachedWriteTokens : 0);
	const prompt = params.prompt ?? [];
	const markerCount = countCacheMarkers(prompt, params.providerOptions);
	activeSpan.setAttribute(SPAN_KEYS.CACHE_INTENT_MARKER_COUNT, markerCount);
	const prefixSignature = computePrefixSignature(prompt, params.tools);
	activeSpan.setAttribute(SPAN_KEYS.CACHE_INTENT_PREFIX_SIGNATURE, prefixSignature);
	const historyMarkerCount = countHistoryCacheMarkers(prompt);
	activeSpan.setAttribute(SPAN_KEYS.CACHE_INTENT_HISTORY_MARKER_COUNT, historyMarkerCount);
}
const gatewayCostMiddleware = {
	specificationVersion: "v3",
	overrideModelId({ model }) {
		return normalizeModelId(model.modelId);
	},
	async wrapGenerate({ doGenerate, params }) {
		const result = await doGenerate();
		try {
			setGatewayAttributesOnSpan(result.providerMetadata);
		} catch (error) {
			logger.warn({ error }, "Failed to extract gateway cost in wrapGenerate");
		}
		try {
			setCacheAttributesOnSpan(params, result.usage);
		} catch (error) {
			logger.warn({ error }, "Failed to set cache attributes in wrapGenerate");
		}
		return result;
	},
	async wrapStream({ doStream, params }) {
		const { stream, ...rest } = await doStream();
		return {
			stream: stream.pipeThrough(new TransformStream({ transform(chunk, controller) {
				controller.enqueue(chunk);
				if (chunk.type === "finish") {
					try {
						setGatewayAttributesOnSpan(chunk.providerMetadata);
					} catch (error) {
						logger.warn({ error }, "Failed to extract gateway cost in wrapStream");
					}
					try {
						setCacheAttributesOnSpan(params, chunk.usage);
					} catch (error) {
						logger.warn({ error }, "Failed to set cache attributes in wrapStream");
					}
				}
			} })),
			...rest
		};
	}
};

//#endregion
export { computePrefixSignature, countCacheMarkers, countHistoryCacheMarkers, extractUsageTokens, gatewayCostMiddleware, normalizeModelId };