import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
import { setTimeout as delay } from "node:timers/promises";

const DEFAULT_HOST = "127.0.0.1";
const DEFAULT_PORT = 11434;
const APFEL_MODEL_ID = "apple-foundationmodel";

type ApfelHealth = {
	status?: string;
	model?: string;
	version?: string;
	active_requests?: number;
	context_window?: number;
	model_available?: boolean;
	prewarmed?: boolean;
	supported_languages?: string[];
};

function configuredHost(): string {
	return process.env.APFEL_HOST || DEFAULT_HOST;
}

function configuredPort(): number {
	const raw = process.env.APFEL_PORT;
	const parsed = raw ? Number(raw) : DEFAULT_PORT;
	return Number.isInteger(parsed) && parsed > 0 ? parsed : DEFAULT_PORT;
}

function configuredBaseUrl(): string {
	return (
		process.env.APFEL_BASE_URL ||
		`http://${configuredHost()}:${configuredPort()}/v1`
	).replace(/\/$/, "");
}

function healthUrl(): string {
	return configuredBaseUrl().replace(/\/v1$/, "") + "/health";
}

async function fetchHealth(timeoutMs = 2_000): Promise<ApfelHealth> {
	const response = await fetch(healthUrl(), {
		signal: AbortSignal.timeout(timeoutMs),
	});
	if (!response.ok) {
		throw new Error(
			`Apfel health returned HTTP ${response.status}: ${await response.text()}`,
		);
	}
	return (await response.json()) as ApfelHealth;
}

async function waitForHealth(timeoutMs: number): Promise<ApfelHealth> {
	const deadline = Date.now() + timeoutMs;
	let lastError: unknown;
	while (Date.now() < deadline) {
		try {
			return await fetchHealth(1_000);
		} catch (error) {
			lastError = error;
			await delay(250);
		}
	}
	throw lastError instanceof Error
		? lastError
		: new Error("Timed out waiting for Apfel health");
}

function formatHealth(health: ApfelHealth): string {
	const languages = health.supported_languages?.length
		? health.supported_languages.join(", ")
		: "unknown";
	return [
		`status: ${health.status ?? "unknown"}`,
		`model: ${health.model ?? APFEL_MODEL_ID}`,
		`version: ${health.version ?? "unknown"}`,
		`available: ${health.model_available ?? "unknown"}`,
		`prewarmed: ${health.prewarmed ?? "unknown"}`,
		`context: ${health.context_window ?? 4096}`,
		`supported languages: ${languages}`,
	].join("\n");
}

export default function apfelProviderExtension(pi: ExtensionAPI) {
	let ownedServer: ChildProcessWithoutNullStreams | undefined;

	pi.registerProvider("apfel", {
		name: "Apfel",
		baseUrl: configuredBaseUrl(),
		api: "openai-completions",
		// Apfel does not require auth by default, but pi uses auth presence when
		// deciding whether a custom local provider is selectable. The value is a
		// harmless placeholder unless the Apfel server was started with --token.
		apiKey: process.env.APFEL_TOKEN || "apfel",
		models: [
			{
				id: APFEL_MODEL_ID,
				name: "Apple Foundation Model via Apfel",
				reasoning: false,
				input: ["text"],
				contextWindow: 4096,
				maxTokens: 4096,
				cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
				compat: {
					supportsDeveloperRole: false,
					supportsReasoningEffort: false,
					maxTokensField: "max_tokens",
				},
			},
		],
	});

	pi.registerCommand("apfel-health", {
		description: "Check the local Apfel provider server",
		handler: async (_args, ctx) => {
			try {
				const health = await fetchHealth();
				ctx.ui.notify(
					formatHealth(health),
					health.model_available === false ? "warning" : "info",
				);
			} catch (error) {
				const message = error instanceof Error ? error.message : String(error);
				ctx.ui.notify(
					`Apfel is not reachable at ${healthUrl()}\n${message}`,
					"error",
				);
			}
		},
	});

	pi.registerCommand("apfel-start", {
		description: "Start `apfel --serve` for the Apfel provider",
		handler: async (args, ctx) => {
			try {
				const health = await fetchHealth(750);
				ctx.ui.notify(
					`Apfel is already running.\n${formatHealth(health)}`,
					"info",
				);
				return;
			} catch {
				// Not running yet; spawn below.
			}

			if (ownedServer && !ownedServer.killed) {
				ctx.ui.notify(
					"Apfel server was already started by this extension.",
					"info",
				);
				return;
			}

			const extraArgs = args.trim() ? args.trim().split(/\s+/) : [];
			const commandArgs = [
				"--serve",
				"--host",
				configuredHost(),
				"--port",
				String(configuredPort()),
				...extraArgs,
			];
			ownedServer = spawn("apfel", commandArgs, {
				env: {
					...process.env,
					APFEL_HOST: configuredHost(),
					APFEL_PORT: String(configuredPort()),
				},
			});

			let stderr = "";
			ownedServer.stderr.on("data", (chunk) => {
				stderr += String(chunk);
				if (stderr.length > 4_000) stderr = stderr.slice(-4_000);
			});
			ownedServer.on("exit", (code, signal) => {
				ownedServer = undefined;
				if (ctx.hasUI)
					ctx.ui.notify(
						`Apfel server exited: ${signal ?? code ?? "unknown"}`,
						"warning",
					);
			});

			try {
				const health = await waitForHealth(10_000);
				ctx.ui.notify(
					`Started Apfel provider at ${configuredBaseUrl()}\n${formatHealth(health)}`,
					"info",
				);
			} catch (error) {
				ownedServer.kill();
				ownedServer = undefined;
				const message = error instanceof Error ? error.message : String(error);
				ctx.ui.notify(
					`Failed to start Apfel. Is it installed and supported on this Mac?\n${message}\n${stderr}`.trim(),
					"error",
				);
			}
		},
	});

	pi.registerCommand("apfel-stop", {
		description: "Stop the Apfel server started by this extension",
		handler: async (_args, ctx) => {
			if (!ownedServer) {
				ctx.ui.notify("No Apfel server is owned by this Pi session.", "info");
				return;
			}
			ownedServer.kill("SIGTERM");
			ownedServer = undefined;
			ctx.ui.notify("Stopped Apfel server.", "info");
		},
	});

	pi.on("session_start", async (_event, ctx) => {
		ctx.ui.setStatus("apfel", `apfel: ${configuredBaseUrl()}`);
	});

	pi.on("session_shutdown", async () => {
		if (ownedServer && !ownedServer.killed) {
			ownedServer.kill("SIGTERM");
		}
		ownedServer = undefined;
	});
}