import { createServer, type IncomingHttpHeaders, type Server } from "node:http";
import type { AddressInfo } from "node:net";
import { afterAll, afterEach, beforeAll, beforeEach, describe, expect, it } from "vitest";
import { getModel, getModels } from "../src/models.js";
import { openCodeSessionHeaders } from "../src/providers/opencode-headers.js";
import { complete, completeSimple, stream } from "../src/stream.js";
import type { Api, AssistantMessage, Context, Model, SimpleStreamOptions } from "../src/types.js";

describe("Cache Retention (PI_CACHE_RETENTION)", () => {
	const originalEnv = process.env.PI_CACHE_RETENTION;

	beforeEach(() => {
		delete process.env.PI_CACHE_RETENTION;
	});

	afterEach(() => {
		if (originalEnv !== undefined) {
			process.env.PI_CACHE_RETENTION = originalEnv;
		} else {
			delete process.env.PI_CACHE_RETENTION;
		}
	});

	const context: Context = {
		systemPrompt: "You are a helpful assistant.",
		messages: [{ role: "user", content: "Hello", timestamp: Date.now() }],
	};

	describe("Anthropic Provider", () => {
		it.skipIf(!process.env.ANTHROPIC_API_KEY)(
			"should use default cache TTL (no ttl field) when PI_CACHE_RETENTION is not set",
			async () => {
				const model = getModel("anthropic", "claude-haiku-4-5");
				let capturedPayload: any = null;

				const s = stream(model, context, {
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const _ of s) {
				}

				expect(capturedPayload).not.toBeNull();
				expect(capturedPayload.system).toBeDefined();
				expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral" });
			},
		);

		it.skipIf(!process.env.ANTHROPIC_API_KEY)("should use 1h cache TTL when PI_CACHE_RETENTION=long", async () => {
			process.env.PI_CACHE_RETENTION = "long";
			const model = getModel("anthropic", "claude-haiku-4-5");
			let capturedPayload: any = null;

			const s = stream(model, context, {
				onPayload: (payload) => {
					capturedPayload = payload;
				},
			});

			for await (const _ of s) {
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.system).toBeDefined();
			expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
		});

		it("should add ttl for non-api.anthropic.com baseUrl by default", async () => {
			process.env.PI_CACHE_RETENTION = "long";

			const baseModel = getModel("anthropic", "claude-haiku-4-5");
			const proxyModel = {
				...baseModel,
				baseUrl: "https://my-proxy.example.com/v1",
			};

			let capturedPayload: any = null;

			const { streamAnthropic } = await import("../src/providers/anthropic.js");

			try {
				const s = streamAnthropic(proxyModel, context, {
					apiKey: "fake-key",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
		});

		it("should omit ttl when supportsLongCacheRetention is false", async () => {
			const baseModel = getModel("anthropic", "claude-haiku-4-5");
			const proxyModel = {
				...baseModel,
				baseUrl: "https://my-proxy.example.com/v1",
				compat: { supportsLongCacheRetention: false },
			};
			let capturedPayload: any = null;

			const { streamAnthropic } = await import("../src/providers/anthropic.js");

			try {
				const s = streamAnthropic(proxyModel, context, {
					apiKey: "fake-key",
					cacheRetention: "long",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral" });
		});

		it("should omit cache_control when cacheRetention is none", async () => {
			const baseModel = getModel("anthropic", "claude-haiku-4-5");
			let capturedPayload: any = null;

			const { streamAnthropic } = await import("../src/providers/anthropic.js");

			try {
				const s = streamAnthropic(baseModel, context, {
					apiKey: "fake-key",
					cacheRetention: "none",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.system[0].cache_control).toBeUndefined();
		});

		it("should add cache_control to string user messages", async () => {
			const baseModel = getModel("anthropic", "claude-haiku-4-5");
			let capturedPayload: any = null;

			const { streamAnthropic } = await import("../src/providers/anthropic.js");

			try {
				const s = streamAnthropic(baseModel, context, {
					apiKey: "fake-key",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			const lastMessage = capturedPayload.messages[capturedPayload.messages.length - 1];
			expect(Array.isArray(lastMessage.content)).toBe(true);
			const lastBlock = lastMessage.content[lastMessage.content.length - 1];
			expect(lastBlock.cache_control).toEqual({ type: "ephemeral" });
		});

		it("should set 1h cache TTL when cacheRetention is long", async () => {
			const baseModel = getModel("anthropic", "claude-haiku-4-5");
			let capturedPayload: any = null;

			const { streamAnthropic } = await import("../src/providers/anthropic.js");

			try {
				const s = streamAnthropic(baseModel, context, {
					apiKey: "fake-key",
					cacheRetention: "long",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.system[0].cache_control).toEqual({ type: "ephemeral", ttl: "1h" });
		});
	});

	describe("OpenAI Responses Provider", () => {
		it.skipIf(!process.env.OPENAI_API_KEY)(
			"should not set prompt_cache_retention when PI_CACHE_RETENTION is not set",
			async () => {
				const model = getModel("openai", "gpt-4o-mini");
				let capturedPayload: any = null;

				const s = stream(model, context, {
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const _ of s) {
				}

				expect(capturedPayload).not.toBeNull();
				expect(capturedPayload.prompt_cache_retention).toBeUndefined();
			},
		);

		it.skipIf(!process.env.OPENAI_API_KEY)(
			"should set prompt_cache_retention to 24h when PI_CACHE_RETENTION=long",
			async () => {
				process.env.PI_CACHE_RETENTION = "long";
				const model = getModel("openai", "gpt-4o-mini");
				let capturedPayload: any = null;

				const s = stream(model, context, {
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const _ of s) {
				}

				expect(capturedPayload).not.toBeNull();
				expect(capturedPayload.prompt_cache_retention).toBe("24h");
			},
		);

		it("should set prompt_cache_retention for non-api.openai.com baseUrl by default", async () => {
			process.env.PI_CACHE_RETENTION = "long";

			const baseModel = getModel("openai", "gpt-4o-mini");
			const proxyModel = {
				...baseModel,
				baseUrl: "https://my-proxy.example.com/v1",
			};

			let capturedPayload: any = null;

			const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");

			try {
				const s = streamOpenAIResponses(proxyModel, context, {
					apiKey: "fake-key",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.prompt_cache_retention).toBe("24h");
		});

		it("should omit prompt_cache_retention when supportsLongCacheRetention is false", async () => {
			const model = {
				...getModel("openai", "gpt-4o-mini"),
				compat: { supportsLongCacheRetention: false },
			};
			let capturedPayload: any = null;

			const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");

			try {
				const s = streamOpenAIResponses(model, context, {
					apiKey: "fake-key",
					cacheRetention: "long",
					sessionId: "session-compat-false",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.prompt_cache_retention).toBeUndefined();
		});

		it("should omit prompt_cache_key when cacheRetention is none", async () => {
			const model = getModel("openai", "gpt-4o-mini");
			let capturedPayload: any = null;

			const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");

			try {
				const s = streamOpenAIResponses(model, context, {
					apiKey: "fake-key",
					cacheRetention: "none",
					sessionId: "session-1",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.prompt_cache_key).toBeUndefined();
			expect(capturedPayload.prompt_cache_retention).toBeUndefined();
		});

		it("should set prompt_cache_retention when cacheRetention is long", async () => {
			const model = getModel("openai", "gpt-4o-mini");
			let capturedPayload: any = null;

			const { streamOpenAIResponses } = await import("../src/providers/openai-responses.js");

			try {
				const s = streamOpenAIResponses(model, context, {
					apiKey: "fake-key",
					cacheRetention: "long",
					sessionId: "session-2",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.prompt_cache_key).toBe("session-2");
			expect(capturedPayload.prompt_cache_retention).toBe("24h");
		});
	});

	describe("OpenAI Completions Provider", () => {
		function createCompletionsModel(compat?: Model<"openai-completions">["compat"]): Model<"openai-completions"> {
			return {
				id: "test-model",
				name: "Test Model",
				api: "openai-completions",
				provider: "test-openai-completions",
				baseUrl: "https://my-proxy.example.com/v1",
				reasoning: false,
				input: ["text"],
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
				contextWindow: 128000,
				maxTokens: 4096,
				compat,
			};
		}

		async function capturePayload(
			model: Model<"openai-completions">,
			options: { cacheRetention?: "none" | "short" | "long"; sessionId?: string },
		) {
			let capturedPayload: any = null;
			const { streamOpenAICompletions } = await import("../src/providers/openai-completions.js");

			try {
				const s = streamOpenAICompletions(model, context, {
					apiKey: "fake-key",
					...options,
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			return capturedPayload;
		}

		it.each([
			{ name: "caching enabled", options: { sessionId: "session-1" }, key: "session-1", retention: undefined },
			{
				name: "cacheRetention long",
				options: { cacheRetention: "long" as const, sessionId: "session-2" },
				key: "session-2",
				retention: "24h",
			},
			{
				name: "cacheRetention none",
				options: { cacheRetention: "none" as const, sessionId: "session-3" },
				key: undefined,
				retention: undefined,
			},
			{
				name: "PI_CACHE_RETENTION=long",
				options: { sessionId: "session-4" },
				env: "long",
				key: "session-4",
				retention: "24h",
			},
		])("should shape direct OpenAI prompt cache fields for $name", async ({ options, env, key, retention }) => {
			if (env) process.env.PI_CACHE_RETENTION = env;
			const { compat: _compat, ...baseModel } = getModel("openai", "gpt-4o-mini");
			const model = { ...baseModel, api: "openai-completions" } as Model<"openai-completions">;

			const payload = await capturePayload(model, options);

			expect(payload.prompt_cache_key).toBe(key);
			expect(payload.prompt_cache_retention).toBe(retention);
		});

		it("should set prompt_cache_retention for non-api.openai.com baseUrl by default", async () => {
			let capturedPayload: any = null;
			const { streamOpenAICompletions } = await import("../src/providers/openai-completions.js");

			try {
				const s = streamOpenAICompletions(createCompletionsModel(), context, {
					apiKey: "fake-key",
					cacheRetention: "long",
					sessionId: "session-completions",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.prompt_cache_key).toBe("session-completions");
			expect(capturedPayload.prompt_cache_retention).toBe("24h");
		});

		it("should omit prompt_cache_retention when supportsLongCacheRetention is false", async () => {
			let capturedPayload: any = null;
			const { streamOpenAICompletions } = await import("../src/providers/openai-completions.js");

			try {
				const s = streamOpenAICompletions(createCompletionsModel({ supportsLongCacheRetention: false }), context, {
					apiKey: "fake-key",
					cacheRetention: "long",
					sessionId: "session-completions-false",
					onPayload: (payload) => {
						capturedPayload = payload;
					},
				});

				for await (const event of s) {
					if (event.type === "error") break;
				}
			} catch {
				// The fake proxy request fails after the payload capture used by this assertion.
			}

			expect(capturedPayload).not.toBeNull();
			expect(capturedPayload.prompt_cache_key).toBeUndefined();
			expect(capturedPayload.prompt_cache_retention).toBeUndefined();
		});
	});
});

const identityContext: Context = { messages: [{ role: "user", content: "Reply OK", timestamp: 1 }] };
const identityModels = (["opencode", "opencode-go"] as const).flatMap((provider) => {
	const catalog = getModels(provider);
	return [...new Set(catalog.map((model) => model.api))].map((api) => catalog.find((model) => model.api === api)!);
});

let server: Server;
let fixtureBaseUrl: string;
let requests: { headers: IncomingHttpHeaders; body: Record<string, unknown> }[];
let enforceContract: boolean;

beforeAll(async () => {
	server = createServer(async (req, res) => {
		let raw = "";
		for await (const chunk of req) raw += chunk.toString();
		const body = JSON.parse(raw) as Record<string, unknown>;
		requests.push({ headers: req.headers, body });
		if (
			enforceContract &&
			(!req.headers["x-opencode-session"] || !req.headers["user-agent"]?.startsWith("prime-agent"))
		) {
			res.writeHead(400, { "content-type": "application/json" });
			res.end(JSON.stringify({ error: { type: "invalid_request_error", message: "Missing OpenCode identity" } }));
			return;
		}

		res.writeHead(200, { "content-type": "text/event-stream" });
		const event = (type: string, data: Record<string, unknown>) =>
			res.write(`event: ${type}\ndata: ${JSON.stringify({ type, ...data })}\n\n`);
		if (req.url?.includes("/messages")) {
			event("message_start", {
				message: {
					id: "fixture-message",
					type: "message",
					role: "assistant",
					model: body.model,
					content: [],
					stop_reason: null,
					stop_sequence: null,
					usage: { input_tokens: 1, output_tokens: 0 },
				},
			});
			event("content_block_start", { index: 0, content_block: { type: "text", text: "" } });
			event("content_block_delta", { index: 0, delta: { type: "text_delta", text: "OK" } });
			event("content_block_stop", { index: 0 });
			event("message_delta", {
				delta: { stop_reason: "end_turn", stop_sequence: null },
				usage: { output_tokens: 1 },
			});
			event("message_stop", {});
		} else if (req.url?.includes("/responses")) {
			const item = {
				id: "fixture-output",
				type: "message",
				status: "completed",
				role: "assistant",
				content: [{ type: "output_text", text: "OK", annotations: [] }],
			};
			event("response.created", { response: { id: "fixture-response", status: "in_progress", output: [] } });
			event("response.output_item.added", {
				output_index: 0,
				item: { ...item, status: "in_progress", content: [] },
			});
			event("response.content_part.added", {
				output_index: 0,
				content_index: 0,
				part: { type: "output_text", text: "", annotations: [] },
			});
			event("response.output_text.delta", { output_index: 0, content_index: 0, delta: "OK" });
			event("response.output_item.done", { output_index: 0, item });
			event("response.completed", {
				response: {
					id: "fixture-response",
					status: "completed",
					output: [item],
					usage: { input_tokens: 1, output_tokens: 1 },
				},
			});
		} else if (req.url?.includes(":streamGenerateContent")) {
			res.write(
				`data: ${JSON.stringify({ candidates: [{ content: { role: "model", parts: [{ text: "OK" }] }, finishReason: "STOP" }] })}\n\n`,
			);
		} else {
			res.write(
				`data: ${JSON.stringify({ id: "fixture-completion", choices: [{ index: 0, delta: { role: "assistant", content: "OK" }, finish_reason: "stop" }] })}\n\n`,
			);
			res.write("data: [DONE]\n\n");
		}
		res.end();
	});
	await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
	fixtureBaseUrl = `http://127.0.0.1:${(server.address() as AddressInfo).port}/v1`;
});

afterAll(async () => {
	server.closeAllConnections();
	await new Promise<void>((resolve, reject) => server.close((error) => (error ? reject(error) : resolve())));
});

beforeEach(() => {
	requests = [];
	enforceContract = true;
});

function expectSuccess(result: AssistantMessage) {
	expect(result.stopReason, result.errorMessage).toBe("stop");
	expect(
		result.content
			.filter((block) => block.type === "text")
			.map((block) => block.text)
			.join(""),
	).toBe("OK");
}

describe.each(identityModels)("request identity headers: $provider/$api", (catalogModel) => {
	function request(options: SimpleStreamOptions = {}, model: Model<Api> = catalogModel, simple = true) {
		return (simple ? completeSimple : complete)({ ...model, baseUrl: fixtureBaseUrl }, identityContext, {
			apiKey: "fixture-not-a-real-key",
			sessionId: "conversation-a",
			maxTokens: 512,
			transport: "sse",
			signal: AbortSignal.timeout(5000),
			...options,
		});
	}

	it.each([true, false])("identifies repeated and distinct conversations (simple=%s)", async (simple) => {
		for (const sessionId of ["conversation-a", "conversation-a", "conversation-b"]) {
			const result = await request({ sessionId }, catalogModel, simple);
			const headers = requests.at(-1)?.headers;
			expect({ session: headers?.["x-opencode-session"], userAgent: headers?.["user-agent"] }).toEqual({
				session: sessionId,
				userAgent: expect.stringMatching(/^prime-agent(?:\/|$)/),
			});
			expectSuccess(result);
		}
	});

	it.each(["none", "short", "long"] as const)(
		"identifies conversations with cacheRetention=%s",
		async (cacheRetention) => {
			const result = await request({ cacheRetention });
			expect(requests.at(-1)?.headers["x-opencode-session"]).toBe("conversation-a");
			expectSuccess(result);
			if (cacheRetention === "none") {
				expect(requests.at(-1)?.body.prompt_cache_key).toBeUndefined();
				expect(requests.at(-1)?.headers.session_id).toBeUndefined();
			}
		},
	);

	it("lets model headers and then request headers override generated identity regardless of casing", async () => {
		const model = {
			...catalogModel,
			headers: {
				"User-Agent": "prime-agent-model/1",
				"X-OpenCode-Session": "model-conversation",
				"X-Fixture": "model",
			},
		};
		expectSuccess(await request({}, model));
		expect(requests.at(-1)?.headers["user-agent"]).toBe("prime-agent-model/1");
		expect(requests.at(-1)?.headers["x-opencode-session"]).toBe("model-conversation");

		expectSuccess(
			await request(
				{
					headers: {
						"user-agent": "prime-agent-request/1",
						"x-opencode-session": "request-conversation",
						"x-fixture": "request",
					},
				},
				model,
			),
		);
		expect(requests.at(-1)?.headers["user-agent"]).toBe("prime-agent-request/1");
		expect(requests.at(-1)?.headers["x-opencode-session"]).toBe("request-conversation");
		expect(requests.at(-1)?.headers["x-fixture"]).toBe("request");
		expect(model.headers["X-Fixture"]).toBe("model");
	});

	it.each([undefined, ""])("does not invent a conversation when sessionId=%s", async (sessionId) => {
		enforceContract = false;
		expectSuccess(await request({ sessionId }));
		expect(requests.at(-1)?.headers["x-opencode-session"]).toBeUndefined();
		expect(requests.at(-1)?.headers["user-agent"]).toMatch(/^prime-agent(?:\/|$)/);
	});

	it("does not add OpenCode identity for unrelated providers", async () => {
		enforceContract = false;
		expectSuccess(await request({}, { ...catalogModel, provider: "fixture-provider" }));
		expect(requests.at(-1)?.headers["x-opencode-session"]).toBeUndefined();
		expect(requests.at(-1)?.headers["user-agent"]).not.toMatch(/^prime-agent(?:\/|$)/);
	});
});

describe("session affinity headers", () => {
	function affinityModel(): Model<"openai-completions"> {
		return {
			id: "test-model",
			name: "Test Model",
			api: "openai-completions",
			provider: "test-openai-completions",
			baseUrl: fixtureBaseUrl,
			reasoning: false,
			input: ["text"],
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
			contextWindow: 128000,
			maxTokens: 4096,
			compat: { sendSessionAffinityHeaders: true },
		};
	}

	const affinityHeaders = ["session_id", "x-client-request-id", "x-session-affinity"] as const;

	it.each([
		{ name: "sends affinity headers when compat enables them", options: {}, expected: "session-affinity" },
		{ name: "omits affinity headers when cacheRetention is none", options: { cacheRetention: "none" as const } },
		{
			name: "lets explicit headers override generated affinity headers",
			options: {
				headers: {
					session_id: "override",
					"x-client-request-id": "override",
					"x-session-affinity": "override",
				},
			},
			expected: "override",
		},
	])("$name", async ({ options, expected }) => {
		enforceContract = false;
		await completeSimple(affinityModel(), identityContext, {
			apiKey: "fixture-not-a-real-key",
			sessionId: "session-affinity",
			transport: "sse",
			signal: AbortSignal.timeout(5000),
			...options,
		});

		const headers = requests.at(-1)?.headers ?? {};
		for (const header of affinityHeaders) {
			expect(headers[header]).toBe(expected);
		}
	});
});

describe("openCodeSessionHeaders", () => {
	it("names the session for OpenCode providers only", () => {
		for (const provider of ["opencode", "opencode-go"]) {
			expect(openCodeSessionHeaders(provider, "session-1")).toEqual({ "x-opencode-session": "session-1" });
			expect(openCodeSessionHeaders(provider, undefined)).toEqual({});
		}
		for (const provider of ["anthropic", "prime-inference", ""]) {
			expect(openCodeSessionHeaders(provider, "session-1")).toEqual({});
		}
	});
});