import { createAssistantMessageEventStream } from "@earendil-works/pi-ai";
import type { Api, AssistantMessage, AssistantMessageEventStream, Model, ProviderStreams, StreamOptions } from "@earendil-works/pi-ai";
import { asRecord } from "./values.ts";
import type { ExplabsModel } from "./types.ts";

export interface GatewayDetails {
	requestId?: string;
	provider?: string;
	routeDepth?: string;
	routeReason?: string;
	zdr?: string;
	replayRepair?: string;
	chargedUsd?: number;
	isByok?: boolean;
	upstreamUsd?: number;
	ignoredParameters?: string[];
	omittedSampling?: string[];
}

export type GatewayMessage = AssistantMessage & { explabs?: GatewayDetails };

function nonnegativeNumber(value: unknown): number | undefined {
	return typeof value === "number" && Number.isFinite(value) && value >= 0 ? value : undefined;
}

export function gatewayHeaders(headers: Record<string, string>): GatewayDetails {
	const lower = Object.fromEntries(Object.entries(headers).map(([key, value]) => [key.toLowerCase(), value]));
	return {
		requestId: lower["x-request-id"],
		provider: lower["x-gateway-provider"],
		routeDepth: lower["x-gateway-route-depth"],
		routeReason: lower["x-gateway-route-reason"],
		zdr: lower["x-gateway-zdr"],
		replayRepair: lower["x-gateway-replay-repair"],
	};
}

export function gatewayEvent(data: unknown): GatewayDetails {
	const event = asRecord(data);
	if (!event) return {};
	// Responses wraps the final reply in response; Messages reports usage on message_delta.
	const response = asRecord(event.response) ?? event;
	const finalUsage = event.type === "response.completed" || event.type === "response.incomplete" ||
		event.type === "response.failed" || event.type === "message_delta" || event.type === undefined;
	const usage = finalUsage ? asRecord(response.usage) : undefined;
	const details = asRecord(usage?.cost_details);
	const ignored = response["x-experiential-ignored-parameters"];
	const chargedUsd = nonnegativeNumber(usage?.cost);
	const upstreamUsd = nonnegativeNumber(details?.upstream_inference_cost);
	return {
		...(typeof response.provider === "string" ? { provider: response.provider } : {}),
		...(chargedUsd !== undefined ? { chargedUsd } : {}),
		...(typeof usage?.is_byok === "boolean" ? { isByok: usage.is_byok } : {}),
		...(upstreamUsd !== undefined ? { upstreamUsd } : {}),
		...(Array.isArray(ignored) ? { ignoredParameters: ignored.filter((item): item is string => typeof item === "string") } : {}),
	};
}

export function rewriteSampling(payload: unknown, model: Model<Api> & Pick<ExplabsModel, "samplingCapabilities">): unknown {
	const record = asRecord(payload);
	const capabilities = model.samplingCapabilities;
	if (!record || !capabilities) return payload;
	const reasoning = asRecord(record.reasoning);
	const thinking = asRecord(record.thinking);
	const outputConfig = asRecord(record.output_config);
	const effort = reasoning?.effort ?? record.reasoning_effort ?? outputConfig?.effort;
	const reasoningEnabled = effort !== undefined ? effort !== "none" && effort !== "off"
		: thinking !== undefined ? thinking.type !== "disabled"
		: model.reasoning && model.thinkingLevelMap?.off === null;
	const samplingConflictsWithReasoning = capabilities.sampling_requires_reasoning_none && reasoningEnabled;
	const next = { ...record };
	function outsideRange(value: unknown, min: number | undefined, max: number | undefined): boolean {
		return typeof value === "number" && ((min !== undefined && value < min) || (max !== undefined && value > max));
	}
	if (capabilities.supports_temperature === false || samplingConflictsWithReasoning ||
		outsideRange(next.temperature, capabilities.minimum_temperature, capabilities.maximum_temperature)) {
		delete next.temperature;
	}
	if (capabilities.supports_top_p === false || samplingConflictsWithReasoning ||
		outsideRange(next.top_p, capabilities.minimum_top_p, capabilities.maximum_top_p)) {
		delete next.top_p;
	}
	return next;
}

export interface GatewayOptions {
	trace?: boolean;
	traceName?: string;
	userId?: string;
	tags?: Record<string, string>;
}

export function readGatewayOptions(env: NodeJS.ProcessEnv = process.env): GatewayOptions {
	let tags: Record<string, string> | undefined;
	if (env.EXPLABS_TAGS) {
		const parsed = asRecord(JSON.parse(env.EXPLABS_TAGS));
		if (!parsed || Object.values(parsed).some((value) => typeof value !== "string")) {
			throw new Error("EXPLABS_TAGS must be a JSON object with string values");
		}
		if (Object.keys(parsed).length > 16 || Buffer.byteLength(JSON.stringify(parsed)) > 8192) {
			throw new Error("EXPLABS_TAGS must contain at most 16 entries and 8 KiB of JSON");
		}
		tags = parsed as Record<string, string>;
	}
	return { trace: env.EXPLABS_TRACE === "1", traceName: env.EXPLABS_TRACE_NAME, userId: env.EXPLABS_USER_ID, tags };
}

export function withGatewayInstrumentation(api: ProviderStreams, config: GatewayOptions = {}): ProviderStreams {
	function instrument<T extends StreamOptions>(options: T | undefined) {
		let details: GatewayDetails = {};
		let omittedSampling: string[] = [];
		const next = {
			...options,
			headers: {
				...(config.trace && options?.sessionId ? { "X-Explabs-Session-Id": options.sessionId } : {}),
				...(config.traceName ? { "X-Explabs-Trace-Name": config.traceName } : {}),
				...(config.userId ? { "X-Explabs-User-Id": config.userId } : {}),
				...(config.tags ? { "X-Explabs-Tags": JSON.stringify(config.tags) } : {}),
				...options?.headers,
			},
			onPayload: async (payload: unknown, model: Model<Api>) => {
				const rewritten = (await options?.onPayload?.(payload, model)) ?? payload;
				const filtered = rewriteSampling(rewritten, model);
				const before = asRecord(rewritten);
				const after = asRecord(filtered);
				omittedSampling = ["temperature", "top_p"].filter((key) => before && key in before && after && !(key in after));
				return filtered;
			},
			onResponse: async (response, model) => {
				// A retry starts a new attempt; never retain another attempt's charge or route.
				details = gatewayHeaders(response.headers);
				await options?.onResponse?.(response, model);
			},
			onProviderStreamEvent: async (event: unknown, model: Model<Api>) => {
				details = { ...details, ...gatewayEvent(event) };
				await options?.onProviderStreamEvent?.(event, model);
			},
		} satisfies StreamOptions;
		return {
			options: next as T,
			attach(message: GatewayMessage) {
				if (Object.values(details).some((value) => value !== undefined) || omittedSampling.length) {
					message.explabs = { ...details, ...(omittedSampling.length ? { omittedSampling: [...omittedSampling] } : {}) };
				}
			},
		};
	}

	function forward(source: AssistantMessageEventStream, attach: (message: GatewayMessage) => void, model: Model<Api>): AssistantMessageEventStream {
		const target = createAssistantMessageEventStream();
		const started = Date.now();
		void (async () => {
			for await (const event of source) {
				if (event.type === "done") attach(event.message);
				else if (event.type === "error") attach(event.error);
				target.push(event);
			}
			const result = await source.result();
			attach(result);
			target.end(result);
		})().catch((error: unknown) => {
			const message: GatewayMessage = {
				role: "assistant", api: model.api, provider: model.provider, model: model.id,
				content: [], timestamp: started, stopReason: "error",
				errorMessage: error instanceof Error ? error.message : String(error),
				usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0,
					cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } },
			};
			attach(message);
			target.push({ type: "error", reason: "error", error: message });
			target.end(message);
		});
		return target;
	}

	return {
		stream(model, context, options) {
			const request = instrument(options);
			return forward(api.stream(model, context, request.options), request.attach, model);
		},
		streamSimple(model, context, options) {
			const request = instrument(options);
			return forward(api.streamSimple(model, context, request.options), request.attach, model);
		},
		fetchDeferred: api.fetchDeferred ? (model, handle, options) => {
			const request = instrument(options);
			return forward(api.fetchDeferred!(model, handle, request.options), request.attach, model);
		} : undefined,
		cancelDeferred: api.cancelDeferred ? (model, handle, options) => api.cancelDeferred!(model, handle, options) : undefined,
	};
}