import assert from "node:assert/strict";
import { describe, it } from "node:test";
import { createAssistantMessageEventStream, normalizeContext, type Api, type Model, type ProviderStreams, type StreamOptions } from "@earendil-works/pi-ai";
import { anthropicMessagesApi } from "@earendil-works/pi-ai/compat";
import { gatewayEvent, gatewayHeaders, readGatewayOptions, rewriteSampling, withGatewayInstrumentation, type GatewayMessage } from "../src/gateway.ts";
import { formatGatewayDetails, gatewayStatus } from "../src/diagnostics.ts";
import { explabsResponsesApi } from "../src/responses.ts";

const model: Model<"openai-responses"> = {
	id: "cheap", name: "Cheap", api: "openai-responses", provider: "explabs", baseUrl: "https://gateway.example/v1",
	input: ["text"], reasoning: true, thinkingLevelMap: { off: "none", low: "low" }, contextWindow: 32000, maxTokens: 100,
	cost: { input: 1, output: 2, cacheRead: 0.1, cacheWrite: 1.25 },
};
const context = normalizeContext({ messages: [{ role: "user", content: "Hi", timestamp: 0 }] });
function responseEvents(usage: Record<string, unknown>) {
	const item = { type: "message", id: "msg_1", role: "assistant", status: "completed", content: [{ type: "output_text", text: "OK", annotations: [] }] };
	return [
		{ type: "response.created", response: { id: "resp_1", status: "in_progress" } },
		{ type: "response.output_item.added", output_index: 0, item: { ...item, content: [] } },
		{ type: "response.output_text.delta", output_index: 0, content_index: 0, delta: "OK" },
		{ type: "response.output_item.done", output_index: 0, item },
		{ type: "response.completed", response: { id: "resp_1", status: "completed", output: [item], provider: "actual-rung",
			"x-experiential-ignored-parameters": ["reasoning->low"], usage } },
	];
}
function sse(events: unknown[]): Response {
	return new Response(events.map(event => `data: ${JSON.stringify(event)}\n\n`).join(""), {
		headers: { "content-type": "text/event-stream", "x-request-id": "req_1", "x-gateway-provider": "first-rung", "x-gateway-route-depth": "1" },
	});
}

describe("gateway accounting and diagnostics", () => {
	it("distinguishes a missing charge from zero and ignores pre-dispatch estimates", () => {
		assert.deepEqual(gatewayEvent({ type: "message_start", usage: { cost: 100 } }), {});
		assert.deepEqual(gatewayEvent({ type: "message_delta", usage: { cost: 0, is_byok: true, cost_details: { upstream_inference_cost: 0.004 } } }), {
			chargedUsd: 0, isByok: true, upstreamUsd: 0.004,
		});
		assert.deepEqual(gatewayEvent({ type: "response.completed", response: { usage: { cost: -1 } } }), {});
		assert.equal(gatewayHeaders({ "X-Request-ID": "req" }).requestId, "req");
		assert.match(gatewayStatus({ chargedUsd: 0, isByok: true, upstreamUsd: 0.004 }), /charged \$0\.000000 · BYOK · upstream \$0\.004000/);
	});

	for (const method of ["stream", "streamSimple"] as const) {
		for (const charged of [undefined, 0, 0.004]) {
			it(`${method} preserves hooks and token estimates with charge ${charged}`, async () => {
				const events = responseEvents({ input_tokens: 10, output_tokens: 2, total_tokens: 12,
					...(charged !== undefined ? { cost: charged, is_byok: charged === 0 } : {}) });
				const observed: unknown[] = [];
				const api = withGatewayInstrumentation(explabsResponsesApi(), { trace: true, tags: { project: "test" } });
				const samplingModel = { ...model, samplingCapabilities: { supports_temperature: false } };
				const stream = api[method](samplingModel, context, {
					apiKey: "xpl_test", sessionId: "session-test", maxRetries: 0,
					onPayload: payload => ({ ...payload as object, metadata: { test: "hook" }, temperature: 0.3 }),
					onResponse: response => { observed.push(response.headers["x-request-id"]); },
					onProviderStreamEvent: event => { observed.push(event); },
					fetch: async (_input, init) => {
						const headers = new Headers(init?.headers);
						assert.equal(headers.get("x-explabs-session-id"), "session-test");
						assert.equal(headers.get("x-explabs-tags"), '{"project":"test"}');
						const payload = JSON.parse(String(init?.body));
						assert.deepEqual(payload.metadata, { test: "hook" });
						assert.equal(payload.temperature, undefined);
						return sse(events);
					},
				});
				const normalized = [];
				for await (const event of stream) normalized.push(event.type);
				const result = await stream.result() as GatewayMessage;
				assert.equal(result.stopReason, "stop");
				assert.equal(result.explabs?.chargedUsd, charged);
				assert.equal(result.explabs?.provider, "actual-rung");
				assert.equal(result.explabs?.requestId, "req_1");
				assert.deepEqual(result.explabs?.ignoredParameters, ["reasoning->low"]);
				assert.deepEqual(result.explabs?.omittedSampling, ["temperature"]);
				assert.ok(Math.abs(result.usage.cost.total - 0.000014) < 1e-15);
				assert.equal(result.usage.totalTokens, 12);
				assert.equal(normalized.at(-1), "done");
				assert.deepEqual(observed, ["req_1", ...events]);
				assert.match(formatGatewayDetails(result), /Pi token-cost estimate/);
			}
			);
		}
	}

	it("resets request details when a new response attempt starts", async () => {
		const stream = (m: Model<Api>, _ctx: unknown, options?: StreamOptions) => {
			const output = createAssistantMessageEventStream();
			void (async () => {
				await options?.onResponse?.({ status: 200, headers: { "x-request-id": "old" } }, m);
				await options?.onProviderStreamEvent?.({ type: "response.completed", response: { usage: { cost: 10 } } }, m);
				await options?.onResponse?.({ status: 200, headers: { "x-request-id": "new" } }, m);
				const result: GatewayMessage = { role: "assistant", api: m.api, provider: m.provider, model: m.id,
					content: [], stopReason: "aborted", timestamp: Date.now(), usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, totalTokens: 0, cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 } } };
				output.push({ type: "error", reason: "aborted", error: result });
				output.end(result);
			})();
			return output;
		};
		const api = withGatewayInstrumentation({ stream, streamSimple: stream } as ProviderStreams);
		const result = await api.streamSimple(model, context).result() as GatewayMessage;
		assert.equal(result.stopReason, "aborted");
		assert.equal(result.explabs?.requestId, "new");
		assert.equal(result.explabs?.chargedUsd, undefined);
	});

	it("captures Messages charges and preserves requested long cache TTL", async () => {
		const m = { ...model, api: "anthropic-messages", baseUrl: "https://gateway.example",
			thinkingLevelMap: { ...model.thinkingLevelMap, xhigh: "xhigh" },
			compat: { forceAdaptiveThinking: true } };
		const events = [
			{ type: "message_start", message: { id: "msg_1", type: "message", role: "assistant", model: "cheap", content: [], stop_reason: null, stop_sequence: null, usage: { input_tokens: 10, output_tokens: 0 } } },
			{ type: "content_block_start", index: 0, content_block: { type: "text", text: "" } },
			{ type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "OK" } },
			{ type: "content_block_stop", index: 0 },
			{ type: "message_delta", delta: { stop_reason: "end_turn", stop_sequence: null }, usage: { input_tokens: 10, output_tokens: 2, cost: 0.001, is_byok: false } },
			{ type: "message_stop" },
		];
		const result = await withGatewayInstrumentation(anthropicMessagesApi()).streamSimple(m, context, {
			apiKey: "xpl_test", cacheRetention: "long", maxRetries: 0, reasoning: "xhigh",
			fetch: async (input, init) => {
				const url = new URL(input instanceof Request ? input.url : String(input));
				assert.equal(url.origin, "https://gateway.example");
				assert.equal(url.pathname, "/v1/messages");
				const body = init?.body ?? (input instanceof Request ? await input.clone().text() : "");
				assert.match(String(body), /"ttl":"1h"/);
				const payload = JSON.parse(String(body));
				assert.equal(payload.thinking.type, "adaptive");
				assert.equal(payload.output_config.effort, "xhigh");
				return new Response(events.map(e => `event: ${e.type}\ndata: ${JSON.stringify(e)}\n\n`).join(""), { headers: { "content-type": "text/event-stream" } });
			},
		}).result() as GatewayMessage;
		assert.equal(result.stopReason, "stop", result.errorMessage);
		assert.equal(result.explabs?.chargedUsd, 0.001);
	});
});

describe("sampling compatibility", () => {
	it("removes unsupported sampling after request hooks, while preserving off-mode sampling", () => {
		const m = { ...model, samplingCapabilities: { supports_temperature: true, supports_top_p: true, sampling_requires_reasoning_none: true } };
		assert.deepEqual(rewriteSampling({ reasoning: { effort: "high" }, temperature: 0.5, top_p: 0.9, model: "cheap" }, m), { reasoning: { effort: "high" }, model: "cheap" });
		assert.deepEqual(rewriteSampling({ reasoning: { effort: "none" }, temperature: 0.5, top_p: 0.9 }, m), { reasoning: { effort: "none" }, temperature: 0.5, top_p: 0.9 });
		assert.deepEqual(rewriteSampling({ thinking: { type: "enabled" }, temperature: 1 }, m), { thinking: { type: "enabled" } });
		assert.deepEqual(rewriteSampling({ temperature: 1, top_p: 0.9 }, { ...model, samplingCapabilities: { supports_temperature: false } }), { top_p: 0.9 });
		assert.deepEqual(rewriteSampling({ temperature: 0.3, top_p: 0.9 }, {
			...model, samplingCapabilities: { minimum_temperature: 1, maximum_temperature: 1, maximum_top_p: 0.8 },
		}), {});
		assert.deepEqual(rewriteSampling({ temperature: 1 }, {
			...model, samplingCapabilities: { minimum_temperature: 1, maximum_temperature: 1 },
		}), { temperature: 1 });
		assert.deepEqual(rewriteSampling({ temperature: 1 }, {
			...m, thinkingLevelMap: { off: null },
		}), {});
		assert.deepEqual(rewriteSampling({ temperature: 1 }, model), { temperature: 1 });
	});
});

describe("trace configuration", () => {
	it("sends no session or tags by default and validates bounded explicit tags", () => {
		assert.deepEqual(readGatewayOptions({}), { trace: false, traceName: undefined, userId: undefined, tags: undefined });
		assert.deepEqual(readGatewayOptions({ EXPLABS_TAGS: '{"team":"dev"}', EXPLABS_TRACE: "1" }).tags, { team: "dev" });
		assert.throws(() => readGatewayOptions({ EXPLABS_TAGS: '{"key":1}' }), /string values/);
		assert.throws(() => readGatewayOptions({ EXPLABS_TAGS: JSON.stringify({ huge: "a".repeat(8192) }) }), /8 KiB/);
	});
});