import type { ClassifierContext, ClassifierQuestion, ImageContent, JsonObject, JsonValue } from "@earendil-works/pi-ai";
function record(value: unknown): value is Record<string, unknown> {
	return typeof value === "object" && value !== null && !Array.isArray(value);
}
function json(value: unknown, depth = 0): value is JsonValue {
	if (depth > 64) return false;
	if (value === null || typeof value === "string" || typeof value === "boolean") return true;
	if (typeof value === "number") return Number.isFinite(value);
	if (Array.isArray(value)) return value.every((item) => json(item, depth + 1));
	return record(value) && Object.values(value).every((item) => json(item, depth + 1));
}
function state(value: unknown): value is JsonObject { return record(value) && json(value); }
function strings(value: unknown): value is Record<string, string> {
	return record(value) && Object.values(value).every((item) => typeof item === "string");
}

/** Match an offered choice case-insensitively, returning its original spelling. */
export function canonicalChoice(allowed: readonly string[], raw: string): string | undefined {
	const trimmed = raw.trim();
	return allowed.find((value) => value.toLowerCase() === trimmed.toLowerCase());
}

export type SaiAnswerRejection = "missing" | "malformed" | "unknown-choice" | "invalid-domain";

/** Normalized classifier answers retain a reported confidence, but never manufacture one. */
export type SaiClassifierAnswer =
	| { type: "choice"; choice: string; probabilities: Record<string, number>; confidence?: number }
	| { type: "score"; score: number; confidence?: number; legend?: Record<string, string> | string[] }
	| { type: "bool"; probability: number; confidence?: number };

export function validSaiProbability(value: unknown): value is number {
	return typeof value === "number" && Number.isFinite(value) && value >= 0 && value <= 1;
}

/** Validate the typed response for each requested question, independently of its siblings. */
export function normalizeClassifierAnswers(
	questions: Record<string, ClassifierQuestion>,
	raw: Record<string, unknown>,
	source: "wire" | "classifier",
): { answers: Record<string, SaiClassifierAnswer>; rejected: Record<string, SaiAnswerRejection> } {
	const answers: Array<[string, SaiClassifierAnswer]> = [];
	const rejected: Array<[string, SaiAnswerRejection]> = [];
	for (const [id, question] of Object.entries(questions)) {
		if (!Object.hasOwn(raw, id)) { rejected.push([id, "missing"]); continue; }
		const answer = raw[id];
		if (!record(answer)) { rejected.push([id, "malformed"]); continue; }
		const expectedType = question.type === "bool" && source === "wire" ? "noul" : question.type;
		if (answer.type !== expectedType) { rejected.push([id, "malformed"]); continue; }
		if (Object.hasOwn(answer, "confidence") && !validSaiProbability(answer.confidence)) {
			rejected.push([id, "invalid-domain"]); continue;
		}
		const confidence = Object.hasOwn(answer, "confidence") ? { confidence: answer.confidence as number } : {};
		if (question.type === "choice") {
			if (typeof answer.choice !== "string" || !record(answer.probabilities)) {
				rejected.push([id, "malformed"]); continue;
			}
			const allowed = Object.keys(question.criteria);
			const choice = canonicalChoice(allowed, answer.choice);
			if (choice === undefined) { rejected.push([id, "unknown-choice"]); continue; }
			const probabilities: Array<[string, number]> = [];
			const seen = new Set<string>();
			let valid = true;
			for (const [key, value] of Object.entries(answer.probabilities)) {
				const canonical = canonicalChoice(allowed, key);
				if (canonical === undefined || seen.has(canonical) || !validSaiProbability(value)) { valid = false; break; }
				seen.add(canonical);
				probabilities.push([canonical, value]);
			}
			if (!valid) { rejected.push([id, "invalid-domain"]); continue; }
			answers.push([id, { type: "choice", choice, probabilities: Object.fromEntries(probabilities), ...confidence }]);
		} else if (question.type === "score") {
			if (typeof answer.score !== "number" || !Number.isFinite(answer.score)) {
				rejected.push([id, "malformed"]); continue;
			}
			if (answer.score < 0 || answer.score > question.criteria.length - 1) {
				rejected.push([id, "invalid-domain"]); continue;
			}
			let legend: Record<string, string> | string[] | undefined;
			if (Object.hasOwn(answer, "legend")) {
				const supplied = answer.legend;
				const entries = Array.isArray(supplied) ? Array.from(supplied.entries(), ([key, value]) => [String(key), value] as const)
					: record(supplied) ? Object.entries(supplied) : undefined;
				if (!entries || entries.some(([key, label]) => !/^(0|[1-9]\d*)$/.test(key) || question.criteria[Number(key)] !== label || typeof label !== "string")) {
					rejected.push([id, "invalid-domain"]); continue;
				}
				legend = Array.isArray(supplied) ? [...supplied] as string[] : Object.fromEntries(entries) as Record<string, string>;
			}
			answers.push([id, { type: "score", score: answer.score, ...confidence, ...(legend === undefined ? {} : { legend }) }]);
		} else {
			const probability = source === "wire" ? answer.noul : answer.probability;
			if (!validSaiProbability(probability)) { rejected.push([id, "invalid-domain"]); continue; }
			answers.push([id, { type: "bool", probability, ...confidence }]);
		}
	}
	return { answers: Object.fromEntries(answers), rejected: Object.fromEntries(rejected) };
}

/** Convert the unchanged public Sai question envelope to upstream's typed classifier context. */
export function toClassifierContext(payload: unknown): ClassifierContext | undefined {
	if (!record(payload) || !state(payload.state) || !record(payload.questions)) return undefined;
	const questions: Record<string, ClassifierQuestion> = Object.create(null);
	for (const [id, question] of Object.entries(payload.questions)) {
		if (!record(question)) return undefined;
		const raw = question.instructions;
		const instructions = typeof raw === "string" ? raw : record(raw) && typeof raw.question === "string"
			? raw.question + (typeof raw.focus === "string" ? `\nFocus: ${raw.focus}` : "") : undefined;
		if (instructions === undefined) return undefined;
		if (question.type === "choice" && strings(question.criteria)) {
			questions[id] = { type: "choice", instructions, criteria: question.criteria };
		} else if (question.type === "score" && Array.isArray(question.criteria) && question.criteria.every((item): item is string => typeof item === "string")) {
			questions[id] = { type: "score", instructions, criteria: question.criteria };
		} else if (question.type === "noul" || question.type === "bool") {
			const criteria = question.criteria;
			const yes = record(criteria) ? criteria.true : undefined;
			const no = record(criteria) ? criteria.false : undefined;
			if (criteria !== undefined && (!strings(criteria) || typeof yes !== "string" || typeof no !== "string")) return undefined;
			questions[id] = { type: "bool", instructions, criteria: typeof yes === "string" && typeof no === "string" ? { true: yes, false: no } : { true: "True", false: "False" } };
		} else return undefined;
	}
	return { state: payload.state, questions };
}
/** Deployed TokenIn accepts string state; preserve typed upstream context via onPayload. */
export function tokenInWirePayload(payload: unknown): unknown {
	if (!record(payload)) throw new Error("TokenIn classifier requires an object payload");
	if (typeof payload.state === "string") return payload;
	if (!state(payload.state)) throw new Error("TokenIn classifier requires JSON state");
	return { ...payload, state: JSON.stringify(payload.state) };
}
export function fitsClassifierRequest(payload: unknown, maxBytes: number): boolean {
	try {
		// Include string-state escaping, not just the smaller in-memory JSON object.
		const wire = record(payload) && "state" in payload ? tokenInWirePayload(payload) : payload;
		const text = JSON.stringify(wire);
		return typeof text === "string" && Number.isFinite(maxBytes) && maxBytes > 0 && Buffer.byteLength(text, "utf8") <= Math.min(maxBytes, 64 * 1024);
	} catch { return false; }
}

/** Inclusive limit on the completed JSON HTTP body, not a token or model-context estimate. */
export const MAX_CLASSIFIER_BODY_BYTES = 64 * 1024;

/**
 * Inspect supported JSON material, including JSON strings such as the gateway's stringified state.
 * Image-looking `image_url` fields/parts, structured image-type snippets and `data:image/` strings
 * are conservatively refused even in prose or code; this is not an arbitrary-encoding decoder.
 * Non-JSON values or excessive nesting cannot
 * be inspected safely. The final serialized body is inspected again on each fetch attempt.
 */
export function inspectClassifierMaterial(value: unknown, depth = 0): "ok" | "media" | "uninspectable" {
	if (depth > 64) return "uninspectable";
	if (value === null || typeof value === "boolean") return "ok";
	if (typeof value === "number") return Number.isFinite(value) ? "ok" : "uninspectable";
	if (typeof value === "string") {
		if (/data:image\//i.test(value)
			|| /["']image_url["']\s*:|\bimage_url\s*:\s*(?:\{|["']?(?:https?:\/\/|data:))/i.test(value)
			|| /["']type["']\s*:\s*["'](?:image|image_url|input_image)["']/i.test(value)) return "media";
		const trimmed = value.trim();
		if (trimmed.startsWith("{") || trimmed.startsWith("[")) {
			try { return inspectClassifierMaterial(JSON.parse(trimmed), depth + 1); } catch { /* ordinary prose or code */ }
		}
		return "ok";
	}
	if (typeof value !== "object") return "uninspectable";
	if (!Array.isArray(value) && Object.getPrototypeOf(value) !== Object.prototype && Object.getPrototypeOf(value) !== null) return "uninspectable";
	if (record(value) && (Object.hasOwn(value, "image_url") || ["image", "image_url", "input_image"].includes(String(value.type)))) return "media";
	for (const item of Object.values(value)) {
		const result = inspectClassifierMaterial(item, depth + 1);
		if (result !== "ok") return result;
	}
	return "ok";
}

// ---------------------------------------------------------------------------
// Shared image gate
// ---------------------------------------------------------------------------

/** The image formats the TokenIn decisions gateway accepts. */
export const SAI_IMAGE_MIME_TYPES = ["image/png", "image/jpeg", "image/webp"] as const;
export type SaiImageMimeType = (typeof SAI_IMAGE_MIME_TYPES)[number];

/** At most this many images may accompany one decisions request. */
export const MAX_CLASSIFIER_IMAGES = 4;

/**
 * Aggregate cap on the encoded (base64) image data in one request. Counted as base64 characters
 * before any decoding, and deliberately separate from the early text budget. This input-hygiene
 * limit does not establish an image token budget; valid images are currently refused as well.
 */
export const MAX_IMAGE_BYTES = 2 * 1024 * 1024;

/** Which rule refused a set of images. Content-free: it names the rule, never the data. */
export type SaiImageRejection = "count" | "size" | "mime" | "base64" | "signature";

/** A validated copy of the image set, or the input-hygiene rule that refused it. */
export type SaiImageValidation =
	| { ok: true; images: ImageContent[] }
	| { ok: false; reason: SaiImageRejection };

function isSaiImageMimeType(value: string): value is SaiImageMimeType {
	return (SAI_IMAGE_MIME_TYPES as readonly string[]).includes(value);
}

/** Base64 as a real encoder writes it: a full quartet, canonical padding, no trailing bits. */
function decodeImageBase64(data: string): Buffer | undefined {
	if (data.length === 0 || data.length % 4 !== 0 || !/^[A-Za-z0-9+/]+={0,2}$/.test(data)) return undefined;
	const bytes = Buffer.from(data, "base64");
	return bytes.toString("base64") === data ? bytes : undefined;
}

/** The leading bytes each accepted format carries; no pixel decoding. */
function imageSignatureMatches(mimeType: SaiImageMimeType, bytes: Buffer): boolean {
	switch (mimeType) {
		case "image/png":
			return bytes.subarray(0, 8).equals(Buffer.from([0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a]));
		case "image/jpeg":
			return bytes[0] === 0xff && bytes[1] === 0xd8 && bytes[2] === 0xff;
		case "image/webp":
			return bytes.subarray(0, 4).toString("latin1") === "RIFF" && bytes.subarray(8, 12).toString("latin1") === "WEBP";
	}
}

/**
 * The one image gate every decisions caller shares. Count and aggregate encoded size are checked
 * before any decoding; each image must then be a supported MIME type, valid base64, and carry that
 * format's signature. A refusal rejects the whole set: nothing is clipped, omitted, or downgraded
 * to a text-only request, and the reason never carries image data or credentials.
 */
export function validateSaiImages(images: readonly ImageContent[]): SaiImageValidation {
	if (images.length > MAX_CLASSIFIER_IMAGES) return { ok: false, reason: "count" };
	let encoded = 0;
	for (const image of images) encoded += image.data.length;
	if (encoded > MAX_IMAGE_BYTES) return { ok: false, reason: "size" };
	const validated: ImageContent[] = [];
	for (const image of images) {
		if (!isSaiImageMimeType(image.mimeType)) return { ok: false, reason: "mime" };
		const bytes = decodeImageBase64(image.data);
		if (!bytes) return { ok: false, reason: "base64" };
		if (!imageSignatureMatches(image.mimeType, bytes)) return { ok: false, reason: "signature" };
		validated.push({ type: "image", data: image.data, mimeType: image.mimeType });
	}
	return { ok: true, images: validated };
}

/** An actionable, content-free transport error for the rule that refused an image set. */
export function saiImageRejectionMessage(reason: SaiImageRejection): string {
	switch (reason) {
		case "count":
			return `TokenIn accepts at most ${MAX_CLASSIFIER_IMAGES} images per request`;
		case "size":
			return `TokenIn images exceed the ${MAX_IMAGE_BYTES / 1024 / 1024} MiB aggregate encoded cap`;
		case "mime":
			return "TokenIn accepts only PNG, JPEG, or WebP images";
		case "base64":
			return "TokenIn image data is not valid base64";
		case "signature":
			return "TokenIn image data does not match its declared MIME type";
	}
}