/**
 * TokenIn's decisions classifiers: Jev and the gateway's other decisions models (Clef, Perplexity
 * Decider, GPT-6 Luna). They all take the same System One request and answer in the same format, so
 * one implementation serves every model; the model id in the request body picks the deployment.
 *
 * Transport: the gateway serves them as chat deployments, so a request is `POST {baseUrl}/chat/completions`
 * with one user message whose content is the JSON System One request `{state, questions}`, and the
 * answer is the System One body (`{answers, usage}`) as the reply's message content. (Its
 * `/v1/systemone` route is not deployed.) Protocol parsing and retries stay upstream-owned: this runs
 * upstream's System One classifier and only changes the envelope through its `onPayload` and `fetch` hooks.
 *
 * Images travel inside the request's `state` array as `{ type: "image_url", image_url: { url: "data:..." } }`
 * parts; the gateway rejects a top-level `images`. Only models that really read them declare image input.
 *
 * Billing: the gateway bills more than the provider's raw charge and puts what the account was charged
 * in the `x-litellm-response-cost` header (0 for a response-cache hit). That is the only cost reported.
 * The `usage.cost` in the body is the provider's raw charge (and 0 for some models), and the catalog
 * price below is a placeholder, so neither is ever presented as cost.
 */
import type {
	ClassifierApi,
	ClassifierContext,
	ClassifierModel,
	ClassifierOptions,
	ClassifierResult,
	ImageContent,
	Usage,
} from "@earendil-works/pi-ai";
import { classify as classifySystemOne } from "@earendil-works/pi-ai/api/typesafe-system-one";
import { fitsClassifierRequest, tokenInWirePayload } from "./classifier-context.ts";

/**
 * `cost` is a required catalog field, NOT a price. Context windows are the ones pi-ai's OpenRouter
 * catalog lists for the same models. `vision` models were checked against the live gateway with
 * solid-colour images: Perplexity Decider and GPT-6 Luna answered every colour correctly, while Clef and
 * Clef Flash accept an image but answer as if they had not seen it (their token counts follow the
 * base64 text), so they are text-only here even though the OpenRouter catalog lists image input.
 */
function tokenInDecisionsModel(
	id: string,
	name: string,
	contextWindow: number,
	vision = false,
): ClassifierModel<"typesafe-system-one"> {
	return {
		type: "classifier", provider: "tokenin", id, name: `${name} (service-priced)`,
		api: "typesafe-system-one", baseUrl: "https://lite.andlet.me/v1", input: vision ? ["text", "image"] : ["text"], contextWindow,
		cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
	};
}

export const TOKENIN_JEV_CLASSIFIER = tokenInDecisionsModel("jev-1.13", "Jev 1.13", 64000);

/** Cloudflare's Clef Flash: 65,536 tokens of context. */
export const TOKENIN_CLEF_FLASH_CLASSIFIER = tokenInDecisionsModel("cloudflare/clef-flash", "Clef Flash", 65536);

/** Every decisions model Token-In serves, Jev first. */
export const TOKENIN_CLASSIFIERS: readonly ClassifierModel<"typesafe-system-one">[] = [
	TOKENIN_JEV_CLASSIFIER,
	tokenInDecisionsModel("cloudflare/clef", "Clef", 65536),
	TOKENIN_CLEF_FLASH_CLASSIFIER,
	tokenInDecisionsModel("perplexity/pplx-decider-v1.1-27b", "Perplexity Decider 1.1 27B", 262144, true),
	tokenInDecisionsModel("openai/gpt-6-luna-decisions", "GPT-6 Luna Decisions", 1050000, true),
];

const COST_HEADER = "x-litellm-response-cost";
const CACHE_HIT_HEADER = "x-litellm-cache-key";
/** The request without its images: state and questions, as before. */
const MAX_TEXT_BYTES = 64 * 1024;
/** ponytail: 2 MiB of base64 in total, the most checked against the live gateway; raise once larger ones are measured. */
export const MAX_IMAGE_BYTES = 2 * 1024 * 1024;

function record(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}

export interface ServicePricedClassifierResult extends ClassifierResult {
	/** Token counts without a monetary claim: the catalog has no price for these calls. */
	unpricedUsage?: Omit<Usage, "cost">;
	/** What the gateway billed for the call; absent when it did not say. */
	reportedCost?: number;
}

function imagePart(image: ImageContent): unknown {
	return { type: "image_url", image_url: { url: `data:${image.mimeType};base64,${image.data}` } };
}

/** The System One request upstream built, as the gateway's chat body, with any images added to its `state`. */
function chatBody(payload: unknown, images: readonly ImageContent[]): unknown {
	const wire = tokenInWirePayload(payload);
	if (!record(wire) || typeof wire.model !== "string") throw new Error("TokenIn classifier requires a model");
	if (!fitsClassifierRequest(wire, MAX_TEXT_BYTES)) throw new Error("TokenIn classifier payload exceeds 64 KiB");
	if (images.reduce((total, image) => total + image.data.length, 0) > MAX_IMAGE_BYTES) {
		throw new Error(`TokenIn classifier images exceed ${MAX_IMAGE_BYTES / 1024 / 1024} MiB`);
	}
	const { model, ...request } = wire;
	const state = images.length > 0 ? [wire.state, ...images.map(imagePart)] : wire.state;
	return { model, messages: [{ role: "user", content: JSON.stringify({ ...request, state }) }] };
}

/** The System One body a chat completion carries as its message content, or undefined. */
function systemOneBody(completion: unknown): Record<string, unknown> | undefined {
	const choices = record(completion) ? completion.choices : undefined;
	const message = Array.isArray(choices) && record(choices[0]) ? choices[0].message : undefined;
	const content = record(message) ? message.content : undefined;
	if (typeof content !== "string") return undefined;
	try {
		const body: unknown = JSON.parse(content);
		return record(body) ? body : undefined;
	} catch {
		return undefined;
	}
}

/** The amount the account was charged for this response, or undefined when the gateway did not say. */
function billedCost(headers: Headers): number | undefined {
	if ((headers.get(CACHE_HIT_HEADER) ?? "").trim() !== "") return 0;
	const reported = headers.get(COST_HEADER);
	if (reported === null || reported.trim() === "") return undefined;
	const cost = Number(reported);
	return Number.isFinite(cost) && cost >= 0 ? cost : undefined;
}

export const classifyTokenIn = async (
	model: ClassifierModel<ClassifierApi>,
	context: ClassifierContext,
	options?: ClassifierOptions,
): Promise<ServicePricedClassifierResult> => {
	let billed: number | undefined;
	let booleanConfidence: Record<string, number> = {};
	let scoreLegend: Record<string, Record<string, string> | string[]> = {};
	const requestFetch = options?.fetch ?? globalThis.fetch;
	// Upstream's System One client rejects images, so a model that reads them gets them through `state`
	// instead. For a text-only model they stay in the context and upstream reports them as unsupported.
	const images = model.input.includes("image") ? (context.images ?? []) : [];
	const { images: _inState, ...textContext } = context;
	const result: ServicePricedClassifierResult = await classifySystemOne(model, images.length > 0 ? textContext : context, {
		...options,
		onPayload: async (payload, requestModel) => {
			const transformed = await options?.onPayload?.(payload, requestModel);
			return chatBody(transformed === undefined ? payload : transformed, images);
		},
		fetch: async (input, init) => {
			// Reset per attempt: a retry must not inherit pricing from another response.
			billed = undefined;
			booleanConfidence = {};
			scoreLegend = {};
			const url = new URL(input instanceof Request ? input.url : String(input));
			url.pathname = url.pathname.replace(/\/systemone$/u, "/chat/completions");
			const response = await requestFetch(url, init);
			if (!response.ok) return response;
			billed = billedCost(response.headers);
			// An unreadable envelope becomes an empty body: upstream reports the missing answers once,
			// instead of this layer failing the attempt and having it retried (and billed) again.
			const body = systemOneBody(await response.json().catch(() => undefined)) ?? {};
			if (record(body.answers)) for (const [id, answer] of Object.entries(body.answers)) {
				if (!record(answer)) continue;
				if (typeof answer.confidence === "number" && Number.isFinite(answer.confidence) && answer.confidence >= 0 && answer.confidence <= 1) booleanConfidence[id] = answer.confidence;
				const legend = answer.legend;
				if (Array.isArray(legend) && legend.every((value) => typeof value === "string")) scoreLegend[id] = legend;
				else if (record(legend)) {
					const labels: Record<string, string> = {};
					for (const [key, label] of Object.entries(legend)) if (/^\d+$/.test(key) && typeof label === "string") labels[key] = label;
					if (Object.keys(labels).length === Object.keys(legend).length) scoreLegend[id] = labels;
				}
			}
			return new Response(JSON.stringify(body), { status: 200, headers: { "content-type": "application/json" } });
		},
	});
	for (const [id, confidence] of Object.entries(booleanConfidence)) {
		const answer = result.answers[id];
		if (answer?.type === "bool") {
			const reported = { ...answer, confidence }; // not a fresh literal: the public type has no `confidence`
			result.answers[id] = reported;
		}
	}
	for (const [id, legend] of Object.entries(scoreLegend)) {
		const answer = result.answers[id];
		if (answer?.type === "score") {
			const reported = { ...answer, legend };
			result.answers[id] = reported;
		}
	}
	if (result.usage) {
		const { cost: _catalogEstimate, ...counts } = result.usage;
		result.unpricedUsage = counts;
		delete result.usage;
	}
	if (billed !== undefined) result.reportedCost = billed;
	return result;
};