import assert from "node:assert/strict";
import { describe, it } from "node:test";
import type { Api, Model } from "@earendil-works/pi-ai";
import {
	recordedSessionModel,
	restoreSessionModel,
	type SessionModelRegistry,
} from "../src/session-model.ts";

type AnyModel = Model<Api>;

const PROVIDER = "explabs";
const MODEL_ID = "deepseek-v4.1-flash";

function model(id = MODEL_ID, provider = PROVIDER): AnyModel {
	return { id, provider } as AnyModel;
}

function modelChange(provider: string, modelId: string): unknown {
	return { type: "model_change", provider, modelId };
}

function assistantMessage(provider: string, id: string): unknown {
	return { type: "message", message: { role: "assistant", provider, model: id } };
}

interface Harness {
	registry: SessionModelRegistry;
	selected: AnyModel[];
	findCalls: string[];
	refreshCalls: string[];
	sleeps: number[];
}

function harness(input: {
	found?: AnyModel | undefined;
	/** Number of auth checks that report "not configured" before it becomes ready. */
	authFailures?: number;
	authThrows?: boolean;
}): Harness {
	const found = "found" in input ? input.found : model();
	let remainingAuthFailures = input.authFailures ?? 0;
	const state: Harness = {
		selected: [],
		findCalls: [],
		refreshCalls: [],
		sleeps: [],
		registry: {} as SessionModelRegistry,
	};
	state.registry = {
		find(provider, modelId) {
			state.findCalls.push(`${provider}/${modelId}`);
			return found;
		},
		hasConfiguredAuth() {
			if (input.authThrows) {
				throw new Error("auth check exploded");
			}
			if (remainingAuthFailures > 0) {
				remainingAuthFailures -= 1;
				return false;
			}
			return true;
		},
		async refresh(provider) {
			state.refreshCalls.push(provider);
		},
	};
	return state;
}

function run(input: {
	entries: unknown[];
	activeModel?: { provider: string; id: string };
	harness: Harness;
	attempts?: number;
}) {
	return restoreSessionModel({
		providerId: PROVIDER,
		entries: input.entries,
		activeModel: input.activeModel,
		registry: input.harness.registry,
		setModel: (selected) => input.harness.selected.push(selected),
		attempts: input.attempts,
		retryDelayMs: 0,
		sleep: async (ms) => {
			input.harness.sleeps.push(ms);
		},
	});
}

describe("recordedSessionModel", () => {
	it("returns the last model_change", () => {
		assert.deepEqual(
			recordedSessionModel([modelChange(PROVIDER, "a"), modelChange(PROVIDER, MODEL_ID)]),
			{ provider: PROVIDER, modelId: MODEL_ID },
		);
	});

	it("falls back to the assistant message provider/model", () => {
		assert.deepEqual(recordedSessionModel([assistantMessage(PROVIDER, MODEL_ID)]), {
			provider: PROVIDER,
			modelId: MODEL_ID,
		});
	});

	it("prefers whichever entry comes last", () => {
		assert.deepEqual(
			recordedSessionModel([assistantMessage(PROVIDER, MODEL_ID), modelChange("other", "x")]),
			{ provider: "other", modelId: "x" },
		);
		assert.deepEqual(
			recordedSessionModel([modelChange("other", "x"), assistantMessage(PROVIDER, MODEL_ID)]),
			{ provider: PROVIDER, modelId: MODEL_ID },
		);
	});

	it("ignores unrelated, malformed, and non-assistant entries", () => {
		assert.equal(
			recordedSessionModel([
				{ type: "session", version: 3 },
				{ type: "thinking_level_change", thinkingLevel: "medium" },
				{ type: "message", message: { role: "user" } },
				{ type: "model_change" },
				{ type: "model_change", provider: PROVIDER },
				null,
				"nope",
			]),
			undefined,
		);
	});
});

describe("restoreSessionModel", () => {
	it("re-selects the recorded model when auth is already configured", async () => {
		const h = harness({});
		assert.equal(
			await run({ entries: [modelChange(PROVIDER, MODEL_ID)], activeModel: model("other", "x"), harness: h }),
			"restored",
		);
		assert.deepEqual(
			h.selected.map((m) => `${m.provider}/${m.id}`),
			[`${PROVIDER}/${MODEL_ID}`],
		);
		assert.deepEqual(h.refreshCalls, []);
	});

	it("refreshes the provider when the availability snapshot is not ready yet", async () => {
		const h = harness({ authFailures: 1 });
		assert.equal(
			await run({ entries: [modelChange(PROVIDER, MODEL_ID)], harness: h }),
			"restored",
		);
		assert.deepEqual(h.refreshCalls, [PROVIDER]);
		assert.equal(h.selected.length, 1);
	});

	it("keeps waiting while the availability snapshot catches up", async () => {
		const h = harness({ authFailures: 5 });
		assert.equal(
			await run({ entries: [modelChange(PROVIDER, MODEL_ID)], harness: h, attempts: 10 }),
			"restored",
		);
		assert.equal(h.sleeps.length, 4);
	});

	it("gives up without selecting once attempts are exhausted", async () => {
		const h = harness({ authFailures: 99 });
		assert.equal(
			await run({ entries: [modelChange(PROVIDER, MODEL_ID)], harness: h, attempts: 3 }),
			"unavailable",
		);
		assert.deepEqual(h.selected, []);
	});

	it("does nothing when the session used another provider", async () => {
		const h = harness({});
		assert.equal(await run({ entries: [modelChange("volcengine-agent-plan", "glm-5.3-flash")], harness: h }), "skipped");
		assert.deepEqual(h.findCalls, []);
		assert.deepEqual(h.selected, []);
	});

	it("does nothing for a session with no recorded model", async () => {
		const h = harness({});
		assert.equal(await run({ entries: [{ type: "session", version: 3 }], harness: h }), "skipped");
		assert.deepEqual(h.findCalls, []);
	});

	it("leaves pi's fallback alone when the slug is gone from the catalog", async () => {
		const h = harness({ found: undefined });
		assert.equal(await run({ entries: [modelChange(PROVIDER, "retired-slug")], harness: h }), "skipped");
		assert.deepEqual(h.selected, []);
		assert.deepEqual(h.refreshCalls, []);
	});

	it("is a no-op when pi already restored the recorded model", async () => {
		const h = harness({});
		assert.equal(
			await run({
				entries: [modelChange(PROVIDER, MODEL_ID)],
				activeModel: { provider: PROVIDER, id: MODEL_ID },
				harness: h,
			}),
			"skipped",
		);
		assert.deepEqual(h.findCalls, []);
		assert.deepEqual(h.selected, []);
	});

	it("still restores when the active provider matches but the model differs", async () => {
		const h = harness({});
		assert.equal(
			await run({
				entries: [modelChange(PROVIDER, MODEL_ID)],
				activeModel: { provider: PROVIDER, id: "some-other-slug" },
				harness: h,
			}),
			"restored",
		);
		assert.equal(h.selected.length, 1);
	});

	it("surfaces registry refresh failures as a retryable state, not a throw", async () => {
		const h = harness({ authFailures: 1 });
		h.registry.refresh = async () => {
			throw new Error("refresh exploded");
		};
		assert.equal(await run({ entries: [modelChange(PROVIDER, MODEL_ID)], harness: h }), "restored");
	});
});