import type { Server } from "node:http";
import {
	type Api,
	type AssistantMessage,
	type AssistantMessageEventStream,
	type Context,
	type Model,
	type ProviderStreamOptions,
	stream,
} from "@singularity-forge/pi-ai";
import type { ModelRegistry } from "@singularity-forge/pi-coding-agent";
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";

type ProxyStreamFn = (
	model: Model<Api>,
	context: Context,
	options?: ProviderStreamOptions,
) => AssistantMessageEventStream;

export interface ProxyServerOptions {
	port: number;
	modelRegistry: Pick<ModelRegistry, "find" | "getAll" | "getApiKey">;
	onLog?: (message: string) => void;
	streamModel?: ProxyStreamFn;
}

interface OpenAiMessage {
	role?: string;
	content?: string | Array<{ type?: string; text?: string }>;
}

interface OpenAiChatBody {
	model?: string;
	messages?: OpenAiMessage[];
	stream?: boolean;
	temperature?: number;
	max_tokens?: number;
}

interface GoogleStreamBody {
	model?: string;
	contents?: Array<{
		role?: string;
		parts?: Array<{ text?: string }>;
	}>;
	systemInstruction?: {
		parts?: Array<{ text?: string }>;
	};
	stream?: boolean;
	temperature?: number;
	generationConfig?: {
		maxOutputTokens?: number;
	};
}

type RouteKind = "openai" | "google";

export class ProxyServer {
	private server: Server | null = null;
	private boundPort: number | null = null;
	private readonly options: ProxyServerOptions;
	private readonly streamModel: ProxyStreamFn;

	constructor(options: ProxyServerOptions) {
		this.options = options;
		this.streamModel = options.streamModel ?? stream;
	}

	isRunning(): boolean {
		return this.server !== null;
	}

	getPort(): number | null {
		return this.boundPort;
	}

	async start(): Promise<void> {
		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<void>((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(): Promise<void> {
		if (!this.server) {
			return;
		}

		const server = this.server;
		this.server = null;
		this.boundPort = null;

		await new Promise<void>((resolve, reject) => {
			server.close((error) => {
				if (error) {
					reject(error);
					return;
				}
				resolve();
			});
		});
	}

	private async handleCompletionRequest(
		req: express.Request,
		res: express.Response,
		routeKind: RouteKind,
	): Promise<void> {
		const body = req.body as OpenAiChatBody | GoogleStreamBody;
		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 as OpenAiChatBody).max_tokens
				: (body as GoogleStreamBody).generationConfig?.maxOutputTokens;

		const context = this.normalizeContext(body, routeKind);
		const options: ProviderStreamOptions = {
			apiKey,
			temperature: body.temperature,
			maxTokens,
			signal: abortController.signal,
		};

		const eventStream = this.streamModel(model, context, options);
		const shouldStream =
			routeKind === "google"
				? (body as GoogleStreamBody).stream !== false
				: (body as OpenAiChatBody).stream === true;

		if (shouldStream) {
			await this.sendStreamingResponse(eventStream, res, routeKind, model);
			return;
		}

		await this.sendBufferedResponse(eventStream, res, routeKind, model);
	}

	private resolveModelReference(
		bodyModel: string | undefined,
		pathModelId: string | undefined,
	): string | undefined {
		return bodyModel ?? pathModelId;
	}

	private resolveModel(modelReference: string): Model<Api> | undefined {
		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);
	}

	private normalizeContext(
		body: OpenAiChatBody | GoogleStreamBody,
		routeKind: RouteKind,
	): Context {
		if (routeKind === "google") {
			return this.normalizeGoogleContext(body as GoogleStreamBody);
		}

		return this.normalizeOpenAiContext(body as OpenAiChatBody);
	}

	private normalizeOpenAiContext(body: OpenAiChatBody): Context {
		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,
		};
	}

	private normalizeGoogleContext(body: GoogleStreamBody): Context {
		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" as const, 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,
		};
	}

	private normalizeOpenAiMessage(
		message: OpenAiMessage,
	): Context["messages"][number] {
		if (message.role === "assistant") {
			return this.createAssistantMessage(
				this.normalizeContent(message.content),
			);
		}

		return this.createUserMessage(this.normalizeContent(message.content));
	}

	private createUserMessage(
		content: string | { type: "text"; text: string }[],
	): Context["messages"][number] {
		return {
			role: "user",
			content,
			timestamp: Date.now(),
		};
	}

	private createAssistantMessage(
		content: string | { type: "text"; text: string }[],
	): AssistantMessage {
		const normalizedContent =
			typeof content === "string"
				? [{ type: "text" as const, text: content }]
				: content;

		return {
			role: "assistant",
			content: normalizedContent,
			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",
			timestamp: Date.now(),
		};
	}

	private normalizeContent(
		content: string | Array<{ type?: string; text?: string }> | undefined,
	): string | { type: "text"; text: string }[] {
		if (typeof content === "string") {
			return content;
		}

		return (content ?? [])
			.filter((part) => typeof part.text === "string")
			.map((part) => ({ type: "text" as const, text: part.text ?? "" }));
	}

	private async sendStreamingResponse(
		eventStream: AssistantMessageEventStream,
		res: express.Response,
		routeKind: RouteKind,
		model: Model<Api>,
	): Promise<void> {
		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();
	}

	private async sendBufferedResponse(
		eventStream: AssistantMessageEventStream,
		res: express.Response,
		routeKind: RouteKind,
		model: Model<Api>,
	): Promise<void> {
		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,
		});
	}

	private extractText(message: AssistantMessage): string {
		return message.content
			.filter(
				(
					content,
				): content is Extract<
					AssistantMessage["content"][number],
					{ type: "text" }
				> => content.type === "text",
			)
			.map((content) => content.text)
			.join("");
	}

	private buildOpenAiChunk(
		model: Model<Api>,
		delta: string,
	): Record<string, unknown> {
		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 }],
		};
	}

	private buildGoogleChunk(delta: string): Record<string, unknown> {
		return {
			candidates: [
				{
					content: {
						parts: [{ text: delta }],
					},
				},
			],
		};
	}
}

export function createProxyServer(options: ProxyServerOptions): ProxyServer {
	return new ProxyServer(options);
}
