import assert from "node:assert/strict";
import { afterEach, describe, it } from 'vitest';
import type {
	Api,
	AssistantMessageEventStream,
	Model,
} from "@singularity-forge/pi-ai";
import { AuthStorage, ModelRegistry } from "@singularity-forge/pi-coding-agent";
import { createProxyServer, type ProxyServer } from "../proxy-server.ts";

let serverCleanup: ProxyServer | undefined;

afterEach(async () => {
	if (serverCleanup) {
		await serverCleanup.stop();
		serverCleanup = undefined;
	}
});

function createFakeStream(): AssistantMessageEventStream {
	const events: Array<
		| { type: "start"; partial: ReturnType<typeof buildAssistantMessage> }
		| {
				type: "text_delta";
				contentIndex: number;
				delta: string;
				partial: ReturnType<typeof buildAssistantMessage>;
		  }
		| {
				type: "done";
				reason: "stop";
				message: ReturnType<typeof buildAssistantMessage>;
		  }
	> = [];
	let finalResult: ReturnType<typeof buildAssistantMessage> | undefined;
	let completed = false;

	const stream = {
		push(event: (typeof events)[number]) {
			events.push(event);
			if (event.type === "done") {
				completed = true;
				finalResult = event.message;
			}
		},
		end(): void {
			completed = true;
			finalResult = buildAssistantMessage([]);
		},
		result(): Promise<ReturnType<typeof buildAssistantMessage>> {
			if (finalResult) {
				return Promise.resolve(finalResult);
			}
			return new Promise((resolve) => {
				const interval = setInterval(() => {
					if (finalResult) {
						clearInterval(interval);
						resolve(finalResult);
					}
				}, 0);
			});
		},
		async *[Symbol.asyncIterator](): AsyncIterator<(typeof events)[number]> {
			let cursor = 0;
			while (!completed || cursor < events.length) {
				const event = events[cursor];
				if (event) {
					cursor++;
					yield event;
					continue;
				}

				await new Promise((resolve) => setTimeout(resolve, 0));
			}
		},
	} as unknown as AssistantMessageEventStream;

	queueMicrotask(() => {
		stream.push({
			type: "start",
			partial: buildAssistantMessage([]),
		});
		stream.push({
			type: "text_delta",
			contentIndex: 0,
			delta: "hello",
			partial: buildAssistantMessage([]),
		});
		stream.push({
			type: "done",
			reason: "stop",
			message: buildAssistantMessage([{ type: "text", text: "hello" }]),
		});
	});
	return stream;
}

function buildAssistantMessage(content: { type: "text"; text: string }[]) {
	return {
		role: "assistant" as const,
		content,
		api: "google-gemini-cli" as Api,
		provider: "google-gemini-cli",
		model: "proxy",
		usage: {
			input: 0,
			output: 0,
			cacheRead: 0,
			cacheWrite: 0,
			totalTokens: 0,
			cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
		},
		stopReason: "stop" as const,
		timestamp: Date.now(),
	};
}

function createRegistry() {
	const authStorage = AuthStorage.inMemory({
		openai: { type: "api_key", key: "sk-test" },
	});
	const modelRegistry = new ModelRegistry(authStorage, undefined);
	return { modelRegistry };
}

function createRegistryWithoutCredentials() {
	return {
		modelRegistry: new ModelRegistry(AuthStorage.inMemory({}), undefined),
	};
}

function createProxyServerForTests(modelRegistry: ModelRegistry): ProxyServer {
	return createProxyServer({
		port: 0,
		modelRegistry: {
			find: (provider, modelId) => modelRegistry.find(provider, modelId),
			getAll: () => modelRegistry.getAll(),
			getApiKey: (model) => modelRegistry.getApiKey(model),
		},
		streamModel: () => createFakeStream(),
	});
}

function findOpenAiModel(modelRegistry: ModelRegistry): Model<Api> {
	const model = modelRegistry
		.getAll()
		.find((candidate) => candidate.provider === "openai");
	if (!model) {
		throw new Error("Expected at least one openai model in the registry");
	}
	return model;
}

describe("ProxyServer", () => {
	it("serves model listings on /v1/models", async () => {
		const { modelRegistry } = createRegistry();
		const server = createProxyServerForTests(modelRegistry);
		serverCleanup = server;
		await server.start();

		const response = await fetch(
			`http://127.0.0.1:${server.getPort()}/v1/models`,
		);
		const data = (await response.json()) as {
			object: string;
			data: Array<{ object: string }>;
		};

		assert.deepEqual(
			{ ok: response.ok, object: data.object, hasModels: data.data.length > 0 },
			{ ok: true, object: "list", hasModels: true },
		);
	});

	it("serves OpenAI completions on /v1/chat/completions", async () => {
		const { modelRegistry } = createRegistry();
		const model = findOpenAiModel(modelRegistry);
		const server = createProxyServerForTests(modelRegistry);
		serverCleanup = server;
		await server.start();

		const response = await fetch(
			`http://127.0.0.1:${server.getPort()}/v1/chat/completions`,
			{
				method: "POST",
				headers: { "content-type": "application/json" },
				body: JSON.stringify({
					model: `${model.provider}/${model.id}`,
					messages: [{ role: "user", content: "hello" }],
				}),
			},
		);
		const data = (await response.json()) as {
			choices: Array<{ message: { content: string } }>;
		};

		assert.deepEqual(data.choices[0].message.content, "hello");
	});

	it("streams Google content on /v1beta/models/:modelId:streamGenerateContent", async () => {
		const { modelRegistry } = createRegistry();
		const model = findOpenAiModel(modelRegistry);
		const server = createProxyServerForTests(modelRegistry);
		serverCleanup = server;
		await server.start();

		const response = await fetch(
			`http://127.0.0.1:${server.getPort()}/v1beta/models/${encodeURIComponent(`${model.provider}/${model.id}`)}:streamGenerateContent`,
			{
				method: "POST",
				headers: { "content-type": "application/json" },
				body: JSON.stringify({
					contents: [{ role: "user", parts: [{ text: "hello" }] }],
				}),
			},
		);
		const text = await response.text();

		assert.deepEqual(text.includes("hello"), true);
	});

	it("returns 401 when credentials are absent", async () => {
		const { modelRegistry } = createRegistryWithoutCredentials();
		const model = modelRegistry
			.getAll()
			.find((candidate) => candidate.provider === "openai");
		if (!model) {
			throw new Error("Expected at least one openai model in the registry");
		}
		const server = createProxyServerForTests(modelRegistry);
		serverCleanup = server;
		await server.start();

		const response = await fetch(
			`http://127.0.0.1:${server.getPort()}/v1/chat/completions`,
			{
				method: "POST",
				headers: { "content-type": "application/json" },
				body: JSON.stringify({
					model: `${model.provider}/${model.id}`,
					messages: [{ role: "user", content: "hello" }],
				}),
			},
		);

		assert.deepEqual(response.status, 401);
	});
});
