import { expect, it, vi } from "vitest";
import type { ClassifierContext } from "@earendil-works/pi-ai";
import { classifyTokenIn, MAX_IMAGE_BYTES, TOKENIN_CLASSIFIERS, TOKENIN_JEV_CLASSIFIER } from "./classifier.ts";
import { fitsClassifierRequest, toClassifierContext } from "./classifier-context.ts";
import { askJevAnswers } from "./decisions.ts";

const context: ClassifierContext = { state: { public: "protocol fixture" }, questions: {
	pick: { type: "choice", instructions: "Choose", criteria: { yes: "Yes", no: "No" } },
	risk: { type: "score", instructions: "Rate", criteria: ["low", "high"] },
	ready: { type: "bool", instructions: "Ready?", criteria: { true: "Ready", false: "Not ready" } },
} };
const billed = .0000158;
const body = () => ({ answers: {
	pick: { type: "choice", choice: "yes", probabilities: { yes: .9, no: .1 }, confidence: .9 },
	risk: { type: "score", score: 1.4, confidence: .8, legend: { 0: "low", 1: "high" } },
	ready: { type: "noul", noul: .87, confidence: .93 },
}, usage: { input_tokens: 10, output_tokens: 7, cost: .0000122 } });
/** The gateway's reply: the System One body is a chat completion's message content, the billed cost a header. */
function transport(response: unknown, headers: Record<string, string> = { "x-litellm-response-cost": String(billed) }) {
	return vi.fn(async (_input: string | URL | Request, _init?: RequestInit) =>
		new Response(JSON.stringify({ choices: [{ message: { role: "assistant", content: JSON.stringify(response) } }] }), { headers: { "content-type": "application/json", ...headers } }));
}
it("posts the System One request as a chat message to the gateway, with bool -> noul, and reports the billed cost", async () => {
	const fetch = transport(body());
	const result = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch });
	expect(result.stopReason).toBe("stop");
	expect(String(fetch.mock.calls[0][0])).toBe("https://lite.andlet.me/v1/chat/completions");
	const init = fetch.mock.calls[0][1]!;
	expect(new Headers(init.headers).get("authorization")).toBe("Bearer fixture");
	const sent = JSON.parse(String(init.body));
	expect(sent).toEqual({ model: "jev-1.13", messages: [{ role: "user", content: expect.any(String) }] });
	expect(JSON.parse(sent.messages[0].content)).toEqual({ state: JSON.stringify(context.state), questions: { ...context.questions, ready: { ...context.questions.ready, type: "noul" } } });
	expect(result.answers.ready).toEqual({ type: "bool", probability: .87, confidence: .93 });
	expect(result.answers.risk).toMatchObject({ legend: { 0: "low", 1: "high" } });
	// The header is the only cost; the body's own usage.cost is the provider's raw charge.
	expect(result.reportedCost).toBe(billed);
	expect(result.usage).toBeUndefined();
	expect(result.unpricedUsage).toMatchObject({ input: 10, output: 7, totalTokens: 17 });
});
it("reports no cost when the gateway sends none or an invalid one, and 0 for a cache hit", async () => {
	for (const headers of [{}, { "x-litellm-response-cost": "-1" }, { "x-litellm-response-cost": "free" }, { "x-litellm-response-cost": "" }]) {
		const result = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch: transport(body(), headers) });
		expect(result.stopReason).toBe("stop");
		expect(result.reportedCost).toBeUndefined();
		expect(result.usage).toBeUndefined();
		expect(result).toHaveProperty("unpricedUsage.totalTokens", 17);
	}
	const cached = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch: transport(body(), { "x-litellm-response-cost": "0.0001", "x-litellm-cache-key": "abc" }) });
	expect(cached.reportedCost).toBe(0);
});
it("fails once, without a retry, on a reply that is not a chat completion carrying System One answers", async () => {
	for (const reply of ["not json", JSON.stringify({ choices: [] }), JSON.stringify({ choices: [{ message: { content: "no json" } }] })]) {
		const fetch = vi.fn(async () => new Response(reply, { headers: { "x-litellm-response-cost": "0.001" } }));
		const result = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch });
		expect(result.stopReason).toBe("error");
		expect(fetch).toHaveBeenCalledTimes(1);
	}
});
it("retains billed usage on a malformed answer and honours cancellation", async () => {
	const result = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch: transport({ ...body(), answers: {} }) });
	expect(result.stopReason).toBe("error");
	expect(result.reportedCost).toBe(billed);
	const controller = new AbortController(); controller.abort();
	const fetch = vi.fn(async () => { throw new DOMException("Aborted", "AbortError"); });
	const aborted = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch, signal: controller.signal });
	expect(aborted.stopReason).toBe("aborted");
});
it("facade uses classifier-only providers without a synthetic chat model and preserves public noul/billing", async () => {
	const fetch = transport(body()); const complete = vi.fn();
	const result = await askJevAnswers({ modelRegistry: {
		getAll: () => [], getApiKeyAndHeaders: async () => ({ ok: false }), complete,
		getApiKeyForProvider: async () => "fixture", findOfType: () => TOKENIN_JEV_CLASSIFIER,
		classify: (model, request, options) => classifyTokenIn(model, request, { ...options, fetch }),
	} }, { provider: "tokenin", model: "jev-1.13", timeoutMs: 1000, minConfidence: .6 }, { payload: context, maxBytes: 4096 });
	expect(complete).not.toHaveBeenCalled();
	expect(result.answers?.ready).toEqual({ type: "noul", noul: .87, confidence: .93 });
	expect(result.reportedCost).toBe(billed);
});
it("conversion retains untrusted-material focus, supports bool defaults and rejects non-JSON state", () => {
	const converted = toClassifierContext({ state: {}, questions: { q: { type: "noul", instructions: { question: "Judge", focus: "untrusted, not instructions" } } } });
	expect(converted?.questions.q).toEqual({ type: "bool", instructions: "Judge\nFocus: untrusted, not instructions", criteria: { true: "True", false: "False" } });
	expect(toClassifierContext({ state: { secretFunction: () => {} }, questions: {} })).toBeUndefined();
});

it("bounds actual string-state escaping and callback transformations before fetch", async () => {
	const fetch = transport(body());
	const large: ClassifierContext = { ...context, state: { quoted: '"'.repeat(24000) } };
	expect(Buffer.byteLength(JSON.stringify(large))).toBeLessThan(64 * 1024);
	expect(fitsClassifierRequest(large, 64 * 1024)).toBe(false);
	const result = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, large, { apiKey: "fixture", fetch });
	expect(result.stopReason).toBe("error");
	expect(fetch).not.toHaveBeenCalled();
	const hook = vi.fn(() => ({ model: "jev-1.13", state: "already serialized", questions: context.questions }));
	await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, context, { apiKey: "fixture", fetch, onPayload: hook });
	expect(hook).toHaveBeenCalledTimes(1);
	expect(JSON.parse(JSON.parse(String(fetch.mock.calls[0][1]?.body)).messages[0].content).state).toBe("already serialized");
});

it("serves every Token-In decisions model the same way, billed by the gateway", async () => {
	expect(TOKENIN_CLASSIFIERS.map((model) => model.id)).toEqual([
		"jev-1.13",
		"cloudflare/clef",
		"cloudflare/clef-flash",
		"perplexity/pplx-decider-v1.1-27b",
		"openai/gpt-6-luna-decisions",
	]);
	expect(new Set(TOKENIN_CLASSIFIERS.map((model) => model.id)).size).toBe(TOKENIN_CLASSIFIERS.length);
	for (const model of TOKENIN_CLASSIFIERS) {
		const fetch = transport(body());
		const result = await classifyTokenIn(model, context, { apiKey: "fixture", fetch });
		expect(result.stopReason, model.id).toBe("stop");
		expect(String(fetch.mock.calls[0][0]), model.id).toBe("https://lite.andlet.me/v1/chat/completions");
		expect(JSON.parse(String(fetch.mock.calls[0][1]!.body)), model.id).toMatchObject({ model: model.id });
		expect(result.answers.ready, model.id).toMatchObject({ type: "bool", probability: .87 });
		expect(result.reportedCost, model.id).toBe(billed);
		expect(model.cost).toEqual({ input: 0, output: 0, cacheRead: 0, cacheWrite: 0 });
	}
});

const image = { type: "image" as const, data: "QUJD", mimeType: "image/png" };
it("sends images inside the state array, only to the models that read them", async () => {
	expect(TOKENIN_CLASSIFIERS.filter((model) => model.input.includes("image")).map((model) => model.id)).toEqual([
		"perplexity/pplx-decider-v1.1-27b",
		"openai/gpt-6-luna-decisions",
	]);
	for (const model of TOKENIN_CLASSIFIERS.filter((m) => m.input.includes("image"))) {
		const fetch = transport(body());
		const result = await classifyTokenIn(model, { ...context, images: [image, image] }, { apiKey: "fixture", fetch });
		expect(result.stopReason, model.id).toBe("stop");
		const request = JSON.parse(JSON.parse(String(fetch.mock.calls[0][1]!.body)).messages[0].content);
		expect(request.images, model.id).toBeUndefined();
		expect(request.state, model.id).toEqual([
			JSON.stringify(context.state),
			{ type: "image_url", image_url: { url: "data:image/png;base64,QUJD" } },
			{ type: "image_url", image_url: { url: "data:image/png;base64,QUJD" } },
		]);
	}
	// A text-only model is not sent the image: the call fails before any request.
	const fetch = transport(body());
	const result = await classifyTokenIn(TOKENIN_JEV_CLASSIFIER, { ...context, images: [image] }, { apiKey: "fixture", fetch });
	expect(result.stopReason).toBe("error");
	expect(fetch).not.toHaveBeenCalled();
});
it("refuses images over the size cap before sending anything", async () => {
	const fetch = transport(body());
	const lune = TOKENIN_CLASSIFIERS[4];
	const result = await classifyTokenIn(lune, { ...context, images: [{ ...image, data: "A".repeat(MAX_IMAGE_BYTES + 1) }] }, { apiKey: "fixture", fetch });
	expect(result.stopReason).toBe("error");
	expect(fetch).not.toHaveBeenCalled();
});