import { describe, expect, test } from "bun:test";
import type { CompletionResult } from "../src/upstream/types.ts";

import { DEFAULT_CONFIG } from "../src/config/defaults.ts";
import { loadConfig } from "../src/config/load.ts";
import type { RouterConfig } from "../src/config/types.ts";
import { classify, classifyTask, pickQualityAxis, scoreHeuristic } from "../src/router/classify.ts";
import { extractFeatures } from "../src/router/features.ts";
import { TIER_ORDER, type Features, type Tier } from "../src/router/types.ts";
import type { Dispatch, DispatchOptions, UpstreamClient } from "../src/upstream/types.ts";
import { parseChatRequest } from "../src/wire/openai/request.ts";
import type { NormRequest } from "../src/wire/types.ts";

const BASE = loadConfig({});

function cfgWith(over: Partial<RouterConfig>): RouterConfig {
	return { ...BASE, ...over };
}

const TOOLS = [
	{
		type: "function",
		function: {
			name: "read",
			description: "Read a file",
			parameters: { type: "object", properties: { path: { type: "string" } } },
		},
	},
];

function req(messages: unknown[], tools: unknown[] | undefined = TOOLS): NormRequest {
	const body: Record<string, unknown> = { model: "auto", messages };
	if (tools !== undefined) body.tools = tools;
	return parseChatRequest(body, new Headers());
}

function featuresFor(messages: unknown[], tools?: unknown[] | undefined): Features {
	const r = req(messages, tools === undefined ? TOOLS : tools);
	return extractFeatures(r, 5000);
}

/** A tool-result continuation at a given autonomous-loop depth, no other signals. */
function contFeatures(toolLoopDepth: number, over: Partial<Features> = {}): Features {
	return {
		promptTokens: 30000,
		newContentTokens: 200,
		turnDepth: toolLoopDepth,
		toolCount: 12,
		toolSchemaBytes: 9783,
		isToolResultContinuation: true,
		toolLoopDepth,
		distinctToolsUsed: 3,
		lastToolFailed: false,
		repeatedToolCall: false,
		circularToolCall: false,
		hasImages: false,
		hasNewImage: false,
		codeBlocks: 0,
		codeBytes: 0,
		looksLikeDiff: false,
		complexityKeywords: [],
		trivialityKeywords: [],
		requestedReasoning: undefined,
		questionCount: 0,
		isTerseInstruction: false,
		...over,
	};
}

const SYSTEM = { role: "system", content: "You are a coding agent." };

/** Upstream double that fails loudly if the adjudicator is consulted. */
function forbiddenUpstream(): UpstreamClient {
	return {
		dispatch(_opts: DispatchOptions): Promise<Dispatch> {
			throw new Error("dispatch must not be called during classification");
		},
		complete(): Promise<CompletionResult> {
			throw new Error("adjudicator must not be called");
		},
		fetchModels(): Promise<unknown[]> {
			return Promise.resolve([]);
		},
		fetchModelsForUser(): Promise<unknown[]> {
			return Promise.resolve([]);
		},
	};
}

function scriptedUpstream(behaviour: () => Promise<CompletionResult>): UpstreamClient {
	return {
		dispatch(_opts: DispatchOptions): Promise<Dispatch> {
			throw new Error("dispatch must not be called during classification");
		},
		complete: behaviour,
		fetchModels(): Promise<unknown[]> {
			return Promise.resolve([]);
		},
		fetchModelsForUser(): Promise<unknown[]> {
			return Promise.resolve([]);
		},
	};
}

const tierIdx = (t: Tier): number => TIER_ORDER.indexOf(t);
describe("scoreHeuristic", () => {
	test("a mechanical tool-result continuation scores cheaper than a fresh architecture question", async () => {
		// The single most valuable signal in agent traffic: most turns are
		// post-tool-result continuations, and they do not need a frontier model.
		const continuation = scoreHeuristic(
			featuresFor([
				SYSTEM,
				{ role: "user", content: "check the version" },
				{
					role: "assistant",
					content: null,
					tool_calls: [{ id: "c1", type: "function", function: { name: "read", arguments: '{"path":"p.json"}' } }],
				},
				{ role: "tool", tool_call_id: "c1", content: '{"version":"1.0.0"}' },
			]),
			BASE,
		);
		const architecture = scoreHeuristic(
			featuresFor([
				SYSTEM,
				{
					role: "user",
					content:
						"Find the root cause of this deadlock, explain the race between the queue drain and shutdown, and redesign the architecture to remove the invariant violation.",
				},
			]),
			BASE,
		);
		expect(continuation.score).toBeLessThan(architecture.score);
		expect(tierIdx(continuation.tier)).toBeLessThan(tierIdx(architecture.tier));
	});

	test("a failing tool result raises the tier above a clean one", async () => {
		const clean = scoreHeuristic(
			featuresFor([
				SYSTEM,
				{ role: "user", content: "build" },
				{
					role: "assistant",
					content: null,
					tool_calls: [{ id: "c1", type: "function", function: { name: "read", arguments: "{}" } }],
				},
				{ role: "tool", tool_call_id: "c1", content: "ok, build succeeded" },
			]),
			BASE,
		);
		const failed = scoreHeuristic(
			featuresFor([
				SYSTEM,
				{ role: "user", content: "build" },
				{
					role: "assistant",
					content: null,
					tool_calls: [{ id: "c1", type: "function", function: { name: "read", arguments: "{}" } }],
				},
				{ role: "tool", tool_call_id: "c1", content: "make: *** [all] Error 2" },
			]),
			BASE,
		);
		expect(failed.score).toBeGreaterThan(clean.score);
	});

	test("a requested high reasoning effort raises the score", async () => {
		const plain = scoreHeuristic(featuresFor([SYSTEM, { role: "user", content: "tidy this up" }]), BASE);
		const thinking = scoreHeuristic(
			extractFeatures(
				parseChatRequest(
					{ model: "auto", tools: TOOLS, reasoning_effort: "high", messages: [SYSTEM, { role: "user", content: "tidy this up" }] },
					new Headers(),
				),
				5000,
			),
			BASE,
		);
		expect(thinking.score).toBeGreaterThan(plain.score);
	});

	test("the reasoning weight is configurable, so a session-wide level can be discounted", async () => {
		// A harness that pins one reasoning level for a whole session turns this
		// "signal" into a constant that lifts every turn's score. Measured live:
		// the level never changed within 111 of 115 conversations, and 64 of 119
		// hard dispatches reached that tier ONLY via the weight — $6.66 billed
		// against $0.16 for the same tokens on the moderate pick.
		//
		// DEFAULT_CONFIG, not BASE: BASE is loadConfig({}), which reads this
		// machine's real config.yml, and this assertion is about shipped values.
		const features = extractFeatures(
			parseChatRequest(
				{ model: "auto", tools: TOOLS, reasoning_effort: "medium", messages: [SYSTEM, { role: "user", content: "tidy this up" }] },
				new Headers(),
			),
			5000,
		);
		const shipped = scoreHeuristic(features, DEFAULT_CONFIG);
		const discounted = scoreHeuristic(features, {
			...DEFAULT_CONFIG,
			classifier: { ...DEFAULT_CONFIG.classifier, reasoningWeights: { ...DEFAULT_CONFIG.classifier.reasoningWeights, medium: 0 } },
		});
		expect(shipped.score - discounted.score).toBeCloseTo(DEFAULT_CONFIG.classifier.reasoningWeights.medium, 5);
		expect(discounted.reasons.some((r) => /requested reasoning/.test(r))).toBe(false);
	});

	test("ships with the historical weights, so enabling a discount is opt-in", async () => {
		expect(DEFAULT_CONFIG.classifier.reasoningWeights).toEqual({ medium: 0.14, high: 0.24, xhigh: 0.3, max: 0.34 });
	});

	test("always produces a bounded score, a real tier, and its reasoning", async () => {
		const c = scoreHeuristic(featuresFor([SYSTEM, { role: "user", content: "hello" }]), BASE);
		expect(c.score).toBeGreaterThanOrEqual(0);
		expect(c.score).toBeLessThanOrEqual(1);
		expect(TIER_ORDER).toContain(c.tier);
		expect(c.source).toBe("heuristic");
		expect(c.reasons.length).toBeGreaterThan(0);
		expect(c.confidence).toBeGreaterThanOrEqual(0);
		expect(c.confidence).toBeLessThanOrEqual(1);
	});

	test("a shallow tool-result continuation stays trivial", async () => {
		const shallow = scoreHeuristic(contFeatures(2), BASE);
		expect(shallow.tier).toBe("trivial");
	});

	test("a sustained autonomous loop climbs out of trivial", async () => {
		// The failure mode this fixes: a long coding loop pinned to the cheapest
		// tier for dozens of turns because agentic complexity never accumulated.
		const shallow = scoreHeuristic(contFeatures(2), BASE);
		const deep = scoreHeuristic(contFeatures(20), BASE);
		expect(tierIdx(deep.tier)).toBeGreaterThan(tierIdx(shallow.tier));
		expect(deep.tier).not.toBe("trivial");
	});

	test("score increases monotonically with loop depth past the agentic threshold", async () => {
		const depths = [4, 6, 8, 10, 15, 20, 30];
		let prev = -1;
		for (const d of depths) {
			const s = scoreHeuristic(contFeatures(d), BASE).score;
			expect(s).toBeGreaterThanOrEqual(prev);
			prev = s;
		}
	});

	test("pure loop depth never reaches hard on its own, however runaway", async () => {
		// A sustained-but-not-runaway loop tops out in moderate: the calibrated
		// ramp ceiling for ordinary deep work.
		const midRange = scoreHeuristic(contFeatures(30), BASE);
		expect(midRange.tier).toBe("moderate");
		// And so does a runaway one. This reverses an earlier cap of 0.70 that let
		// raw depth buy `hard`: on live data 152 of 155 hard dispatches were
		// depth-driven and carried 63.5% of ALL spend, while those same rows also
		// scored the mechanical tool-result-continuation penalty. `hard` has no
		// price ceiling, so depth alone was buying a ~8x model for work the
		// classifier already knew was mechanical. Depth is a weak signal of
		// DIFFICULTY; hard must be bought by a corroborating stuck signal (see the
		// circular-tool-call test below), which is the case that actually needs a
		// stronger model.
		const runaway = scoreHeuristic(contFeatures(90), BASE);
		expect(runaway.tier).toBe("moderate");
		// The ceiling must still be a real ceiling, not an accident of the ramp.
		expect(scoreHeuristic(contFeatures(400), BASE).tier).toBe("moderate");
	});

	test("a circular tool call on a FRESH turn escalates to hard", async () => {
		// Off a continuation the stuck signal keeps full weight: the user is
		// watching a live loop and a pricier model may actually break it.
		const deepCircular = scoreHeuristic(
			{ ...contFeatures(90, { circularToolCall: true }), isToolResultContinuation: false },
			BASE,
		);
		expect(deepCircular.tier).toBe("hard");
	});

	test("a circular tool call on a mechanical continuation is damped, not hard", async () => {
		// Measured: hard escalations on circular calls never shortened the loop
		// (chain means identical, 5.74 turns, hard vs moderate). 22 of 27 such
		// hard turns were mechanical continuations paying up to 6x for nothing.
		const retry = scoreHeuristic(contFeatures(90, { circularToolCall: true }), BASE);
		const plain = scoreHeuristic(contFeatures(90), BASE);
		expect(tierIdx(retry.tier)).toBeLessThanOrEqual(tierIdx("moderate"));
		// Exactly the damped weight, not merely "less than the full one".
		expect(retry.score - plain.score).toBeCloseTo(BASE.classifier.mechanicalRetryFactor * 0.24, 5);
	});

	test("a failing tool result on a deep loop is at least simple", async () => {
		// Was 'at least moderate' before the mechanical-retry damp: the flat
		// +0.26 pushed deep mechanical retry loops into hard. A damped retry
		// still clears trivial.
		const deepAndFailing = scoreHeuristic(contFeatures(20, { lastToolFailed: true }), BASE);
		expect(tierIdx(deepAndFailing.tier)).toBeGreaterThanOrEqual(tierIdx("simple"));
	});

	test("a failed-tool retry on a mechanical continuation is damped, not hard", async () => {
		// A retry after a failed tool call is the most mechanical turn there is;
		// the flat +0.26 let automated retry loops buy the hard tier ($7.02 of one
		// measured day vs $0.19 for the same rows as moderate picks). The
		// continuation keeps only mechanicalRetryFactor of the weight.
		const retry = scoreHeuristic(contFeatures(20, { lastToolFailed: true }), BASE);
		const quiet = scoreHeuristic(contFeatures(20), BASE);
		expect(tierIdx(retry.tier)).toBeLessThan(tierIdx("hard"));
		expect(retry.score - quiet.score).toBeCloseTo(
			BASE.classifier.mechanicalRetryFactor * 0.26,
			5,
		);
		// A failure the USER sees (not a tool-result continuation) keeps the full
		// weight: that genuinely changes what the turn needs.
		const userSeen = scoreHeuristic(contFeatures(2, { isToolResultContinuation: false, lastToolFailed: true }), BASE);
		expect(userSeen.score - scoreHeuristic(contFeatures(2, { isToolResultContinuation: false }), BASE).score).toBeCloseTo(0.26, 5);
	});
});

describe("pickQualityAxis", () => {
	test("tools imply the coding axis, plain chat the chat axis", async () => {
		expect(pickQualityAxis(featuresFor([SYSTEM, { role: "user", content: "fix it" }]), BASE)).toBe(BASE.classifier.toolAxis);
		expect(pickQualityAxis(featuresFor([SYSTEM, { role: "user", content: "hello" }], []), BASE)).toBe(
			BASE.classifier.chatAxis,
		);
	});

	test("a deep tool loop switches to the agentic axis", async () => {
		const deep = featuresFor([
			SYSTEM,
			{ role: "user", content: "go" },
			{ role: "assistant", content: null, tool_calls: [{ id: "c1", type: "function", function: { name: "read", arguments: '{"path":"a"}' } }] },
			{ role: "tool", tool_call_id: "c1", content: "a" },
			{ role: "assistant", content: null, tool_calls: [{ id: "c2", type: "function", function: { name: "read", arguments: '{"path":"b"}' } }] },
			{ role: "tool", tool_call_id: "c2", content: "b" },
			{ role: "assistant", content: null, tool_calls: [{ id: "c3", type: "function", function: { name: "read", arguments: '{"path":"c"}' } }] },
			{ role: "tool", tool_call_id: "c3", content: "c" },
			{ role: "assistant", content: null, tool_calls: [{ id: "c4", type: "function", function: { name: "read", arguments: '{"path":"d"}' } }] },
			{ role: "tool", tool_call_id: "c4", content: "d" },
		]);
		expect(deep.toolLoopDepth).toBeGreaterThanOrEqual(BASE.classifier.agenticLoopDepth);
		expect(pickQualityAxis(deep, BASE)).toBe("agentic");
	});
});

describe("classify", () => {
	const messages = [SYSTEM, { role: "user", content: "tidy the retry helper a bit" }];

	test("never consults the adjudicator when the ambiguity threshold is zero", async () => {
		const cfg = cfgWith({ classifier: { ...BASE.classifier, ambiguityThreshold: 0 } });
		const r = req(messages);
		const result = await classify(r, extractFeatures(r, 5000), cfg, {
			upstream: forbiddenUpstream(),
			ledger: null,
			catalog: null,
		});
		expect(result.source).toBe("heuristic");
	});

	test("falls back to the heuristic when the adjudicator returns garbage", async () => {
		// Always-ambiguous, so the adjudicator is definitely consulted.
		const cfg = cfgWith({ classifier: { ...BASE.classifier, ambiguityThreshold: 1.1 } });
		const r = req(messages);
		const f = extractFeatures(r, 5000);
		const expected = scoreHeuristic(f, cfg);
		const result = await classify(r, f, cfg, {
			upstream: scriptedUpstream(() => Promise.resolve({ text: "definitely not a tier", costUsd: 0, toolCalls: [] })),
			ledger: null,
			catalog: null,
		});
		expect(result.tier).toBe(expected.tier);
	});

	test("falls back to the heuristic when the adjudicator throws", async () => {
		const cfg = cfgWith({ classifier: { ...BASE.classifier, ambiguityThreshold: 1.1 } });
		const r = req(messages);
		const f = extractFeatures(r, 5000);
		const expected = scoreHeuristic(f, cfg);
		const result = await classify(r, f, cfg, {
			upstream: scriptedUpstream(() => Promise.reject(new Error("upstream exploded"))),
			ledger: null,
			catalog: null,
		});
		expect(result.tier).toBe(expected.tier);
		expect(result.reasons.some((x) => x.toLowerCase().includes("adjudicat"))).toBe(true);
	});

	test("adopts a valid adjudicator verdict", async () => {
		const cfg = cfgWith({ classifier: { ...BASE.classifier, ambiguityThreshold: 1.1 } });
		const r = req(messages);
		const f = extractFeatures(r, 5000);
		const result = await classify(r, f, cfg, {
			upstream: scriptedUpstream(() => Promise.resolve({ text: "hard", costUsd: 0.00001, toolCalls: [] })),
			ledger: null,
			catalog: null,
		});
		expect(result.tier).toBe("hard");
		expect(result.source).toBe("llm");
	});
});

describe("classifyTask", () => {
	test("image input is a vision task", async () => {
		const f = featuresFor([SYSTEM, { role: "user", content: [{ type: "image_url", image_url: { url: "data:image/png;base64,xxx" } }] }], []);
		expect(classifyTask(f)).toBe("vision");
	});

	test("a stale image on a tool continuation is coding, not vision", async () => {
		const f = featuresFor(
			[
				SYSTEM,
				{
					role: "user",
					content: [
						{ type: "text", text: "build this UI" },
						{ type: "image_url", image_url: { url: "data:image/png;base64,xxx" } },
					],
				},
				{ role: "assistant", content: null, tool_calls: [{ id: "c1", type: "function", function: { name: "read", arguments: "{}" } }] },
				{ role: "tool", tool_call_id: "c1", name: "read", content: "ok" },
			],
			TOOLS,
		);
		expect(f.hasImages).toBe(true);
		expect(f.hasNewImage).toBe(false);
		expect(classifyTask(f)).toBe("coding");
	});

	test("a freshly supplied image mid-loop is vision", async () => {
		const f = featuresFor(
			[
				SYSTEM,
				{ role: "user", content: "start" },
				{ role: "assistant", content: null, tool_calls: [{ id: "c1", type: "function", function: { name: "read", arguments: "{}" } }] },
				{ role: "tool", tool_call_id: "c1", name: "read", content: "ok" },
				{
					role: "user",
					content: [
						{ type: "text", text: "here is the error" },
						{ type: "image_url", image_url: { url: "data:image/png;base64,xxx" } },
					],
				},
			],
			TOOLS,
		);
		expect(f.hasNewImage).toBe(true);
		expect(classifyTask(f)).toBe("vision");
	});

	test("code blocks and diffs are coding tasks", async () => {
		expect(classifyTask(featuresFor([SYSTEM, { role: "user", content: "```ts\nconst x = 1;\n```" }], []))).toBe("coding");
		expect(classifyTask(featuresFor([SYSTEM, { role: "user", content: "diff --git a/x b/x\n@@ -1 +1 @@\n-old\n+new" }], []))).toBe("coding");
	});

	test("tools offered is a coding task", async () => {
		expect(classifyTask(featuresFor([SYSTEM, { role: "user", content: "read the file" }], TOOLS))).toBe("coding");
	});

	test("bare chat with no tools or code is a chat task", async () => {
		expect(classifyTask(featuresFor([SYSTEM, { role: "user", content: "hello, how are you?" }], []))).toBe("chat");
	});

	test("design/architecture prose is a documentation task", async () => {
		expect(classifyTask(featuresFor([SYSTEM, { role: "user", content: "explain the architecture of the system" }], []))).toBe("documentation");
	});
});

describe("classifier.readOnlyToolWeight", () => {
	test("subtracts only when enabled and the tail is a read-only loop", async () => {
		const base = { ...featuresFor([{ role: "user", content: "look" }]), isToolResultContinuation: true, readOnlyToolTail: true };
		const off = scoreHeuristic(base, DEFAULT_CONFIG);
		const cfg = structuredClone(DEFAULT_CONFIG);
		cfg.classifier.readOnlyToolWeight = 0.1;
		const on = scoreHeuristic(base, cfg);
		expect(on.score).toBeCloseTo(Math.max(0, off.score - 0.1), 6);
		expect(on.reasons.some((r) => r.includes("read-only tool loop"))).toBe(true);
		expect(scoreHeuristic({ ...base, readOnlyToolTail: false }, cfg).score).toBeCloseTo(off.score, 6);
	});
});