import type {
	ExtensionAPI,
	ExtensionCommandContext,
	ExtensionContext,
} from "@earendil-works/pi-coding-agent";
import type { TUI } from "@earendil-works/pi-tui";
import { CONFIG_PATH, loadConfig, normalizeConfig, saveConfig } from "./config.js";
import { UniversalHudStatusline } from "./render.js";
import type { HudConfig, HudState, LineLayout, ToolRun } from "./types.js";
import { extractToolTarget, isAgentTool, numberValue } from "./utils.js";

/**
 * Universal HUD for Pi, modeled after claude-hud's statusline renderer.
 *
 * Commands:
 *   /uhud                         toggle
 *   /uhud-mode expanded|compact   switch line layout
 *   /uhud-mode statusline|footer  switch render surface
 *   /uhud-config                  show current config
 */

const HUD_WIDGET_ID = "universal-hud";
const TOGGLE_SHORTCUT = "ctrl+shift+u";
const MODE_SHORTCUT = "ctrl+alt+u";
const REFRESH_MS = 700;

const state: HudState = {
	sessionStartedAt: Date.now(),
	assistantTurnActive: false,
	turns: 0,
	activeTools: new Map(),
	toolCounts: new Map(),
	failedToolCounts: new Map(),
	agentCompleted: 0,
	agentFailed: 0,
};

let config: HudConfig = loadConfig();
let currentCtx: ExtensionContext | null = null;
let widgetTui: TUI | null = null;
let footerTui: TUI | null = null;
let widgetOwned = false;
let footerOwned = false;
let refreshTimer: ReturnType<typeof setInterval> | null = null;

export default function (pi: ExtensionAPI) {
	const requestRender = () => {
		widgetTui?.requestRender();
		footerTui?.requestRender();
	};

	const clearWidget = (ctx: ExtensionContext) => {
		if (!widgetOwned) return;
		ctx.ui.setWidget(HUD_WIDGET_ID, undefined);
		widgetOwned = false;
		widgetTui = null;
	};

	const clearFooter = (ctx: ExtensionContext) => {
		if (!footerOwned) return;
		ctx.ui.setFooter(undefined);
		footerOwned = false;
		footerTui = null;
	};

	const showStatusline = (ctx: ExtensionContext) => {
		clearFooter(ctx);
		ctx.ui.setWidget(
			HUD_WIDGET_ID,
			(tui, theme) => {
				widgetTui = tui;
				return new UniversalHudStatusline(pi, () => currentCtx, () => config, state, theme, "statusline");
			},
			{ placement: "belowEditor" },
		);
		widgetOwned = true;
	};

	const showFooter = (ctx: ExtensionContext) => {
		clearWidget(ctx);
		ctx.ui.setFooter((tui, theme) => {
			footerTui = tui;
			return new UniversalHudStatusline(pi, () => currentCtx, () => config, state, theme, "footer");
		});
		footerOwned = true;
	};

	const applyHud = (ctx: ExtensionContext) => {
		if (!ctx.hasUI) return;
		currentCtx = ctx;
		if (!config.enabled) {
			clearWidget(ctx);
			clearFooter(ctx);
			return;
		}
		if (config.surface === "footer") showFooter(ctx);
		else showStatusline(ctx);
		requestRender();
	};

	const updateConfig = (ctx: ExtensionContext, patch: Partial<HudConfig>) => {
		config = normalizeConfig({ ...config, ...patch });
		saveConfig(config);
		applyHud(ctx);
	};

	const startRefreshTimer = () => {
		if (refreshTimer !== null) return;
		refreshTimer = setInterval(() => requestRender(), REFRESH_MS);
	};

	const stopRefreshTimer = () => {
		if (refreshTimer === null) return;
		clearInterval(refreshTimer);
		refreshTimer = null;
	};

	pi.on("session_start", (_event, ctx) => {
		currentCtx = ctx;
		state.sessionStartedAt = Date.now();
		applyHud(ctx);
		startRefreshTimer();
	});

	pi.on("session_shutdown", (_event, ctx) => {
		if (ctx) {
			clearWidget(ctx);
			clearFooter(ctx);
		}
		currentCtx = null;
		state.agentStartedAt = undefined;
		state.turnStartedAt = undefined;
		state.assistantTurnActive = false;
		state.activeTools.clear();
		stopRefreshTimer();
	});

	pi.on("agent_start", () => {
		state.agentStartedAt = Date.now();
		requestRender();
	});

	pi.on("agent_end", () => {
		state.agentStartedAt = undefined;
		requestRender();
	});

	pi.on("turn_start", () => {
		state.assistantTurnActive = true;
		state.turnStartedAt = Date.now();
		state.turns += 1;
		requestRender();
	});

	pi.on("turn_end", (event: any) => {
		state.assistantTurnActive = false;
		const elapsedMs = state.turnStartedAt ? Date.now() - state.turnStartedAt : 0;
		const outputTokens = numberValue(event?.message?.usage?.output, event?.message?.usage?.outputTokens);
		if (elapsedMs > 0 && outputTokens > 0) {
			state.lastOutputSpeed = outputTokens / (elapsedMs / 1000);
		}
		state.turnStartedAt = undefined;
		requestRender();
	});

	pi.on("tool_execution_start", (event) => {
		const run: ToolRun = {
			id: event.toolCallId,
			name: event.toolName,
			target: extractToolTarget(event.toolName, event.args),
			startedAt: Date.now(),
			isAgent: isAgentTool(event.toolName),
		};
		state.activeTools.set(event.toolCallId, run);
		requestRender();
	});

	pi.on("tool_execution_end", (event) => {
		const run = state.activeTools.get(event.toolCallId);
		state.activeTools.delete(event.toolCallId);
		const elapsedMs = run ? Date.now() - run.startedAt : 0;
		state.lastTool = { name: event.toolName, ok: !event.isError, elapsedMs };
		const targetMap = event.isError ? state.failedToolCounts : state.toolCounts;
		targetMap.set(event.toolName, (targetMap.get(event.toolName) ?? 0) + 1);
		if (run?.isAgent) {
			if (event.isError) state.agentFailed += 1;
			else state.agentCompleted += 1;
		}
		requestRender();
	});

	pi.on("model_select", () => requestRender());
	pi.on("thinking_level_select", () => requestRender());
	pi.on("message_update", () => requestRender());
	pi.on("message_end", () => requestRender());

	pi.registerCommand("uhud", {
		description: "Toggle Universal HUD statusline",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			updateConfig(ctx, { enabled: !config.enabled });
			ctx.ui.notify(`Universal HUD ${config.enabled ? "enabled" : "disabled"}.`, "info");
		},
	});

	pi.registerCommand("universal-hud", {
		description: "Toggle Universal HUD statusline",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			updateConfig(ctx, { enabled: !config.enabled });
			ctx.ui.notify(`Universal HUD ${config.enabled ? "enabled" : "disabled"}.`, "info");
		},
	});

	pi.registerCommand("uhud-mode", {
		description: "Set Universal HUD layout or surface",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			const value = args.trim().toLowerCase();
			if (value === "expanded" || value === "compact") {
				updateConfig(ctx, { enabled: true, lineLayout: value });
				ctx.ui.notify(`Universal HUD layout: ${value}.`, "info");
				return;
			}
			if (value === "statusline" || value === "widget") {
				updateConfig(ctx, { enabled: true, surface: "statusline" });
				ctx.ui.notify("Universal HUD surface: statusline below editor.", "info");
				return;
			}
			if (value === "footer") {
				updateConfig(ctx, { enabled: true, surface: "footer" });
				ctx.ui.notify("Universal HUD surface: footer replacement.", "info");
				return;
			}
			if (value.length === 0) {
				const nextLayout: LineLayout = config.lineLayout === "expanded" ? "compact" : "expanded";
				updateConfig(ctx, { enabled: true, lineLayout: nextLayout });
				ctx.ui.notify(`Universal HUD layout: ${nextLayout}.`, "info");
				return;
			}
			ctx.ui.notify("Usage: /uhud-mode [expanded|compact|statusline|footer]", "warning");
		},
	});

	pi.registerCommand("uhud-layout", {
		description: "Set Universal HUD layout: expanded or compact",
		handler: async (args: string, ctx: ExtensionCommandContext) => {
			const value = args.trim().toLowerCase();
			if (value !== "expanded" && value !== "compact") {
				ctx.ui.notify("Usage: /uhud-layout [expanded|compact]", "warning");
				return;
			}
			updateConfig(ctx, { enabled: true, lineLayout: value });
		},
	});

	pi.registerCommand("uhud-config", {
		description: "Show Universal HUD config path and current config",
		handler: async (_args: string, ctx: ExtensionCommandContext) => {
			ctx.ui.notify(`${CONFIG_PATH}\n${JSON.stringify(config, null, 2)}`, "info");
		},
	});

	pi.registerShortcut(TOGGLE_SHORTCUT, {
		description: "Toggle Universal HUD",
		handler: (ctx: ExtensionContext) => updateConfig(ctx, { enabled: !config.enabled }),
	});

	pi.registerShortcut(MODE_SHORTCUT, {
		description: "Switch Universal HUD expanded/compact layout",
		handler: (ctx: ExtensionContext) => {
			const nextLayout: LineLayout = config.lineLayout === "expanded" ? "compact" : "expanded";
			updateConfig(ctx, { enabled: true, lineLayout: nextLayout });
		},
	});
}