import type { ExtensionAPI, ExtensionCommandContext, ExtensionStartupContext } from "@singularity-forge/pi-coding-agent";
import { createProxyServer, type ProxyServer } from "./proxy-server.js";

const PROXY_COMMAND_NAME = "genai-proxy";
const PROXY_FLAG_NAME = "gemini-cli-proxy";
const DEFAULT_PROXY_PORT = 3000;

export interface ProxyCommandDependencies {
	createProxyServer?: typeof createProxyServer;
}

export function installGenaiProxyExtension(
	api: Pick<ExtensionAPI, "registerCommand" | "registerFlag">,
	dependencies?: ProxyCommandDependencies,
): void {
	let proxyServer: ProxyServer | null = null;
	const buildProxyServer = dependencies?.createProxyServer ?? createProxyServer;

	const ensureProxyServer = (context: ExtensionStartupContext | ExtensionCommandContext, port: number): ProxyServer => {
		if (proxyServer && proxyServer.getPort() === port) {
			return proxyServer;
		}
		if (proxyServer) {
			throw new Error(`Proxy already running on port ${proxyServer.getPort()}`);
		}

		proxyServer = buildProxyServer({
			port,
			modelRegistry: context.modelRegistry,
			onLog: (message) => notifyProxyStatus(context, message, "info"),
		});
		return proxyServer;
	};

	api.registerFlag(PROXY_FLAG_NAME, {
		description: "Start the Gemini CLI proxy server",
		type: "string",
		allowNoValue: true,
		onStartup: async (value, context) => {
			const server = ensureProxyServer(context, resolveProxyPort(value));
			await server.start();
		},
	});

	api.registerCommand(PROXY_COMMAND_NAME, {
		description: "Manage the Gemini CLI proxy server",
		handler: async (args, context) => {
			await handleProxyCommand(args ?? "", context, ensureProxyServer, () => proxyServer, () => {
				proxyServer = null;
			});
		},
	});
}

export function resolveProxyPort(flagValue: boolean | string | undefined): number {
	if (flagValue === true || flagValue === false || flagValue === undefined) {
		return DEFAULT_PROXY_PORT;
	}

	const port = Number.parseInt(flagValue, 10);
	if (!Number.isFinite(port) || port <= 0 || port > 65535) {
		throw new Error(`Invalid proxy port: ${flagValue}`);
	}
	return port;
}

async function handleProxyCommand(
	rawArgs: string,
	context: ExtensionCommandContext,
	ensureProxyServer: (context: ExtensionCommandContext, port: number) => ProxyServer,
	getProxyServer: () => ProxyServer | null,
	clearProxyServer: () => void,
): Promise<void> {
	const [subcommand = "status", portArg] = rawArgs
		.trim()
		.split(/\s+/)
		.filter((value): value is string => value.length > 0);

	if (subcommand === "start") {
		const existingServer = getProxyServer();
		if (existingServer?.isRunning()) {
			notifyProxyStatus(context, `Proxy already running on port ${existingServer.getPort()}`, "info");
			return;
		}

		const server = ensureProxyServer(context, resolveProxyPort(portArg === undefined ? true : portArg));
		await server.start();
		return;
	}

	if (subcommand === "stop") {
		const server = getProxyServer();
		if (!server?.isRunning()) {
			notifyProxyStatus(context, "Proxy is not running", "warning");
			return;
		}

		await server.stop();
		clearProxyServer();
		notifyProxyStatus(context, "Proxy stopped", "success");
		return;
	}

	if (subcommand === "status") {
		const server = getProxyServer();
		if (server?.isRunning()) {
			notifyProxyStatus(context, `Proxy running on port ${server.getPort()}`, "info");
			return;
		}

		notifyProxyStatus(context, "Proxy is not running", "info");
		return;
	}

	notifyProxyStatus(context, "Usage: /genai-proxy start [port] | stop | status", "warning");
}

function notifyProxyStatus(
	context: ExtensionStartupContext | ExtensionCommandContext,
	message: string,
	type: Parameters<ExtensionCommandContext["ui"]["notify"]>[1],
): void {
	if ("ui" in context) {
		context.ui.notify(message, type);
		return;
	}

	process.stderr.write(`[genai-proxy] ${message}\n`);
}
