import { stream } from "@singularity-forge/pi-ai";
import express from "express";

const LISTEN_ADDRESS = "127.0.0.1";
const OPENAI_CREATED_TIMESTAMP = 1_677_610_602;
const SSE_CONTENT_TYPE = "text/event-stream";
const NDJSON_CONTENT_TYPE = "application/x-ndjson";
export class ProxyServer {
	server = null;
	boundPort = null;
	options;
	streamModel;
	constructor(options) {
		this.options = options;
		this.streamModel = options.streamModel ?? stream;
	}
	isRunning() {
		return this.server !== null;
	}
	getPort() {
		return this.boundPort;
	}
	async start() {
		if (this.server) {
			return;
		}
		const app = express();
		app.use(express.json({ limit: "2mb" }));
		app.get(["/v1/models", "/v1beta/models"], (_req, res) => {
			const models = this.options.modelRegistry.getAll().map((model) => ({
				id: model.id,
				object: "model",
				created: OPENAI_CREATED_TIMESTAMP,
				owned_by: model.provider,
				name: model.name,
				capabilities: model.capabilities,
			}));
			if (_req.path.startsWith("/v1beta")) {
				res.json({ models });
				return;
			}
			res.json({ object: "list", data: models });
		});
		app.post("/v1/chat/completions", async (req, res) => {
			await this.handleCompletionRequest(req, res, "openai");
		});
		app.post(
			"/v1beta/models/:modelId\\:streamGenerateContent",
			async (req, res) => {
				await this.handleCompletionRequest(req, res, "google");
			},
		);
		await new Promise((resolve, reject) => {
			const server = app.listen(this.options.port, LISTEN_ADDRESS, () => {
				this.server = server;
				const address = server.address();
				if (typeof address === "object" && address) {
					this.boundPort = address.port;
				} else {
					this.boundPort = this.options.port;
				}
				this.options.onLog?.(
					`Proxy Server running on http://${LISTEN_ADDRESS}:${this.boundPort}`,
				);
				resolve();
			});
			server.once("error", reject);
		});
	}
	async stop() {
		if (!this.server) {
			return;
		}
		const server = this.server;
		this.server = null;
		this.boundPort = null;
		await new Promise((resolve, reject) => {
			server.close((error) => {
				if (error) {
					reject(error);
					return;
				}
				resolve();
			});
		});
	}
	async handleCompletionRequest(req, res, routeKind) {
		const body = req.body;
		const modelReference = this.resolveModelReference(
			body.model,
			req.params.modelId,
		);
		if (!modelReference) {
			res.status(400).json({ error: "Model ID is required" });
			return;
		}
		const model = this.resolveModel(modelReference);
		if (!model) {
			res.status(404).json({ error: `Model ${modelReference} not found` });
			return;
		}
		const apiKey = await this.options.modelRegistry.getApiKey(model);
		if (!apiKey) {
			res
				.status(401)
				.json({ error: `No credentials for provider ${model.provider}` });
			return;
		}
		const abortController = new AbortController();
		req.once("close", () => abortController.abort());
		const maxTokens =
			routeKind === "openai"
				? body.max_tokens
				: body.generationConfig?.maxOutputTokens;
		const context = this.normalizeContext(body, routeKind);
		const options = {
			apiKey,
			temperature: body.temperature,
			maxTokens,
			signal: abortController.signal,
		};
		const eventStream = this.streamModel(model, context, options);
		const shouldStream =
			routeKind === "google" ? body.stream !== false : body.stream === true;
		if (shouldStream) {
			await this.sendStreamingResponse(eventStream, res, routeKind, model);
			return;
		}
		await this.sendBufferedResponse(eventStream, res, routeKind, model);
	}
	resolveModelReference(bodyModel, pathModelId) {
		return bodyModel ?? pathModelId;
	}
	resolveModel(modelReference) {
		const normalizedReference = modelReference.toLowerCase();
		const exact = this.options.modelRegistry
			.getAll()
			.find(
				(model) =>
					`${model.provider}/${model.id}`.toLowerCase() ===
						normalizedReference ||
					model.id.toLowerCase() === normalizedReference,
			);
		if (exact) {
			return exact;
		}
		const slashIndex = modelReference.indexOf("/");
		if (slashIndex === -1) {
			return undefined;
		}
		const provider = modelReference.slice(0, slashIndex);
		const modelId = modelReference.slice(slashIndex + 1);
		return this.options.modelRegistry.find(provider, modelId);
	}
	normalizeContext(body, routeKind) {
		if (routeKind === "google") {
			return this.normalizeGoogleContext(body);
		}
		return this.normalizeOpenAiContext(body);
	}
	normalizeOpenAiContext(body) {
		const messages = body.messages ?? [];
		const systemPrompt = messages.find(
			(message) => message.role === "system",
		)?.content;
		const normalizedMessages = messages
			.filter((message) => message.role !== "system")
			.map((message) => this.normalizeOpenAiMessage(message));
		return {
			systemPrompt: typeof systemPrompt === "string" ? systemPrompt : undefined,
			messages: normalizedMessages,
		};
	}
	normalizeGoogleContext(body) {
		const systemPrompt =
			body.systemInstruction?.parts?.map((part) => part.text ?? "").join("") ||
			undefined;
		const normalizedMessages = (body.contents ?? [])
			.map((content) => {
				const textContent = (content.parts ?? [])
					.filter((part) => typeof part.text === "string")
					.map((part) => ({ type: "text", text: part.text ?? "" }));
				if (content.role === "user") {
					return this.createUserMessage(textContent);
				}
				return this.createAssistantMessage(textContent);
			})
			.filter((message) => message.content.length > 0);
		return {
			systemPrompt,
			messages: normalizedMessages,
		};
	}
	normalizeOpenAiMessage(message) {
		if (message.role === "assistant") {
			return this.createAssistantMessage(
				this.normalizeContent(message.content),
			);
		}
		return this.createUserMessage(this.normalizeContent(message.content));
	}
	createUserMessage(content) {
		return {
			role: "user",
			content,
			timestamp: Date.now(),
		};
	}
	createAssistantMessage(content) {
		const normalizedContent =
			typeof content === "string" ? [{ type: "text", text: content }] : content;
		return {
			role: "assistant",
			content: normalizedContent,
			api: "google-gemini-cli",
			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",
			timestamp: Date.now(),
		};
	}
	normalizeContent(content) {
		if (typeof content === "string") {
			return content;
		}
		return (content ?? [])
			.filter((part) => typeof part.text === "string")
			.map((part) => ({ type: "text", text: part.text ?? "" }));
	}
	async sendStreamingResponse(eventStream, res, routeKind, model) {
		res.status(200);
		res.setHeader(
			"Content-Type",
			routeKind === "openai" ? SSE_CONTENT_TYPE : NDJSON_CONTENT_TYPE,
		);
		res.setHeader("Cache-Control", "no-cache");
		res.setHeader("Connection", "keep-alive");
		for await (const event of eventStream) {
			if (event.type === "text_delta") {
				if (routeKind === "openai") {
					res.write(
						`data: ${JSON.stringify(this.buildOpenAiChunk(model, event.delta))}\n\n`,
					);
				} else {
					res.write(`${JSON.stringify(this.buildGoogleChunk(event.delta))}\n`);
				}
			}
			if (event.type === "done") {
				if (routeKind === "openai") {
					res.write("data: [DONE]\n\n");
				}
				res.end();
				return;
			}
			if (event.type === "error") {
				if (!res.headersSent) {
					res
						.status(500)
						.json({ error: event.error.errorMessage ?? "Proxy stream failed" });
				} else {
					res.end();
				}
				return;
			}
		}
		res.end();
	}
	async sendBufferedResponse(eventStream, res, routeKind, model) {
		const assistantMessage = await eventStream.result();
		const text = this.extractText(assistantMessage);
		if (routeKind === "openai") {
			res.json({
				id: `chatcmpl-${Date.now()}`,
				object: "chat.completion",
				created: Math.floor(Date.now() / 1000),
				model: model.id,
				choices: [
					{
						index: 0,
						message: { role: "assistant", content: text },
						finish_reason: "stop",
					},
				],
				usage: assistantMessage.usage,
			});
			return;
		}
		res.json({
			candidates: [
				{
					content: {
						parts: [{ text }],
					},
				},
			],
			usageMetadata: assistantMessage.usage,
		});
	}
	extractText(message) {
		return message.content
			.filter((content) => content.type === "text")
			.map((content) => content.text)
			.join("");
	}
	buildOpenAiChunk(model, delta) {
		return {
			id: `chatcmpl-${Date.now()}`,
			object: "chat.completion.chunk",
			created: Math.floor(Date.now() / 1000),
			model: model.id,
			choices: [{ index: 0, delta: { content: delta }, finish_reason: null }],
		};
	}
	buildGoogleChunk(delta) {
		return {
			candidates: [
				{
					content: {
						parts: [{ text: delta }],
					},
				},
			],
		};
	}
}
export function createProxyServer(options) {
	return new ProxyServer(options);
}
