/**
 * Model Router — LLM-powered automatic model selection
 *
 * Uses Haiku (~200ms, ~$0.001/call) to classify every prompt's complexity,
 * then routes to the optimal model + thinking level. No regex guessing.
 *
 * Tiers:
 *   ⚡ Haiku   — Simple Q&A, file reads, small edits, quick lookups
 *   ⚖️ Sonnet  — Code generation, debugging, refactoring, multi-file work
 *   🧠 Opus    — Architecture, complex reasoning, large-scale changes
 *
 * Mid-turn escalation:
 *   - 4+ edits/writes → escalate to Opus
 *   - 2+ consecutive errors → escalate (model struggling)
 *   - 6+ tool calls in one turn → escalate to at least Sonnet
 *
 * Usage: pi -e extensions/model-router.ts
 *
 * Commands:
 *   /router          — Show routing state & stats
 *   /router lock     — Lock current model (disable auto-routing)
 *   /router unlock   — Re-enable auto-routing
 *   /router tier <n> — Force tier: 1=haiku, 2=sonnet, 3=opus
 */

import type { AssistantMessage } from "@mariozechner/pi-ai";
import { complete } from "@mariozechner/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { truncateToWidth, visibleWidth } from "@mariozechner/pi-tui";
import { basename } from "node:path";
import { applyExtensionDefaults } from "./themeMap.ts";

// ── Model Tiers ────────────────────────────────────────────────────────────

interface ModelTier {
	name: string;
	icon: string;
	provider: string;
	modelId: string;
	thinking: "off" | "low" | "medium" | "high";
	inputCost: number;
	outputCost: number;
}

const TIERS: ModelTier[] = [
	{ name: "Haiku", icon: "⚡", provider: "anthropic", modelId: "claude-haiku-4-5", thinking: "off", inputCost: 0.80, outputCost: 4.00 },
	{ name: "Sonnet", icon: "⚖️", provider: "anthropic", modelId: "claude-sonnet-4-6", thinking: "low", inputCost: 3.00, outputCost: 15.00 },
	{ name: "Opus", icon: "🧠", provider: "anthropic", modelId: "claude-opus-4-6", thinking: "high", inputCost: 15.00, outputCost: 75.00 },
];

// ── LLM Classifier ────────────────────────────────────────────────────────

const CLASSIFIER_PROMPT = `You are a task complexity classifier for a coding assistant. Given a user prompt, respond with ONLY a JSON object (no markdown, no explanation):

{"tier": <1|2|3>, "reason": "<5 words max>"}

TIER 1 (simple): Quick questions, yes/no, read/list/check files, small edits, conversational ("thanks", "ok"), simple lookups.
TIER 2 (medium): Write functions/classes, fix bugs, refactor single files, write tests, create endpoints, docker/CI config.
TIER 3 (complex): Architecture design, multi-file refactoring/migration, security audits, performance optimization, complex algorithms, system design from scratch.`;

interface Classification { tier: number; reason: string }

const cache = new Map<string, { result: Classification; ts: number }>();
const CACHE_TTL = 60_000;

async function classify(prompt: string, ctx: ExtensionContext): Promise<Classification> {
	const key = prompt.trim().toLowerCase().slice(0, 200);
	const c = cache.get(key);
	if (c && Date.now() - c.ts < CACHE_TTL) return c.result;

	const model = ctx.modelRegistry.find("anthropic", "claude-haiku-4-5");
	if (!model) return { tier: 2, reason: "no classifier" };
	const apiKey = await ctx.modelRegistry.getApiKey(model);
	if (!apiKey) return { tier: 2, reason: "no key" };

	try {
		const resp = await complete(model, {
			systemPrompt: CLASSIFIER_PROMPT,
			messages: [{ role: "user" as const, content: prompt, timestamp: Date.now() }],
		}, { reasoning: "off" });

		const text = resp.content
			.filter((c): c is { type: "text"; text: string } => c.type === "text")
			.map(c => c.text).join("");

		const json = text.replace(/```json?\n?/g, "").replace(/```/g, "").trim();
		const parsed = JSON.parse(json);
		const result: Classification = {
			tier: Math.max(1, Math.min(3, parsed.tier || 2)),
			reason: String(parsed.reason || "classified").slice(0, 50),
		};

		cache.set(key, { result, ts: Date.now() });
		return result;
	} catch {
		return { tier: 2, reason: "classify failed" };
	}
}

// ── Router State ───────────────────────────────────────────────────────────

interface RouterState {
	currentTier: number;
	locked: boolean;
	totalSwitches: number;
	turnsSinceLast: number;
	consecutiveErrors: number;
	toolCallsThisTurn: number;
	editWriteThisTurn: number;
	totalToolCalls: number;
	totalErrors: number;
	tierHistory: Array<{ tier: number; reason: string; turn: number }>;
	turnCount: number;
	savedVsAlwaysOpus: number;
	classifierCost: number;
}

// ── Extension ──────────────────────────────────────────────────────────────

export default function modelRouter(pi: ExtensionAPI) {
	const state: RouterState = {
		currentTier: 1,
		locked: false,
		totalSwitches: 0,
		turnsSinceLast: 0,
		consecutiveErrors: 0,
		toolCallsThisTurn: 0,
		editWriteThisTurn: 0,
		totalToolCalls: 0,
		totalErrors: 0,
		tierHistory: [],
		turnCount: 0,
		savedVsAlwaysOpus: 0,
		classifierCost: 0,
	};

	async function switchToTier(tierIndex: number, reason: string, ctx?: ExtensionContext) {
		const clamped = Math.max(0, Math.min(2, tierIndex));
		if (clamped === state.currentTier) return;
		const tier = TIERS[clamped];
		const model = ctx?.modelRegistry.find(tier.provider, tier.modelId);
		if (!model) return;
		const success = await pi.setModel(model);
		if (!success) return;
		pi.setThinkingLevel(tier.thinking);
		const prev = state.currentTier;
		state.currentTier = clamped;
		state.totalSwitches++;
		state.turnsSinceLast = 0;
		state.tierHistory.push({
			tier: clamped,
			reason: `${TIERS[prev].icon}→${tier.icon} ${reason}`,
			turn: state.turnCount,
		});
	}

	// ── LLM Classification on every prompt ─────────────────────────────────

	pi.on("before_agent_start", async (event, ctx) => {
		if (state.locked) return;

		ctx.ui.setWorkingMessage("classifying...");
		const result = await classify(event.prompt, ctx);
		ctx.ui.setWorkingMessage();

		// Track classifier cost (~180 tokens per call at Haiku rates)
		state.classifierCost += (150 * 0.80 + 30 * 4.00) / 1_000_000;

		let target = result.tier - 1; // 1-3 → 0-2

		// Context-aware overrides
		if (target === 0 && state.totalToolCalls > 10) target = 1;
		if (state.consecutiveErrors >= 2 && target < 2) target = Math.min(target + 1, 2);

		if (target !== state.currentTier) {
			await switchToTier(target, result.reason, ctx);
		}
	});

	// ── Turn & Tool Tracking ───────────────────────────────────────────────

	pi.on("turn_start", async () => {
		state.toolCallsThisTurn = 0;
		state.editWriteThisTurn = 0;
		state.turnCount++;
		state.turnsSinceLast++;
	});

	pi.on("tool_execution_start", async (event) => {
		state.toolCallsThisTurn++;
		state.totalToolCalls++;
		if (event.toolName === "edit" || event.toolName === "write") state.editWriteThisTurn++;
	});

	pi.on("tool_execution_end", async (event, ctx) => {
		if (event.isError) {
			state.consecutiveErrors++;
			state.totalErrors++;
			if (state.consecutiveErrors >= 2 && !state.locked && state.currentTier < 2) {
				await switchToTier(state.currentTier + 1, `${state.consecutiveErrors} errors`, ctx);
			}
		} else {
			state.consecutiveErrors = 0;
		}
	});

	pi.on("turn_end", async (_event, ctx) => {
		if (!state.locked) {
			if (state.editWriteThisTurn >= 4 && state.currentTier < 2) {
				await switchToTier(2, `heavy: ${state.editWriteThisTurn} edits`, ctx);
			} else if (state.toolCallsThisTurn >= 6 && state.currentTier < 1) {
				await switchToTier(1, `busy: ${state.toolCallsThisTurn} tools`, ctx);
			}
		}
		const opusCost = (2000 * TIERS[2].inputCost + 1000 * TIERS[2].outputCost) / 1e6;
		const actualCost = (2000 * TIERS[state.currentTier].inputCost + 1000 * TIERS[state.currentTier].outputCost) / 1e6;
		state.savedVsAlwaysOpus += opusCost - actualCost;
	});

	// ── Session Start ──────────────────────────────────────────────────────

	pi.on("session_start", async (_event, ctx) => {
		applyExtensionDefaults(import.meta.url, ctx);

		const initialTier = TIERS[1];
		const model = ctx.modelRegistry.find(initialTier.provider, initialTier.modelId);
		if (model) {
			await pi.setModel(model);
			pi.setThinkingLevel(initialTier.thinking);
		}

		ctx.ui.setFooter((tui, theme, footerData) => {
			const unsub = footerData.onBranchChange(() => tui.requestRender());
			return {
				dispose: unsub,
				invalidate() {},
				render(width: number): string[] {
					const tier = TIERS[state.currentTier];
					const lockIcon = state.locked ? " 🔒" : "";

					let tokIn = 0, tokOut = 0, cost = 0;
					for (const entry of ctx.sessionManager.getBranch()) {
						if (entry.type === "message" && entry.message.role === "assistant") {
							const m = entry.message as AssistantMessage;
							tokIn += m.usage.input;
							tokOut += m.usage.output;
							cost += m.usage.cost.total;
						}
					}
					cost += state.classifierCost;
					const fmt = (n: number) => n < 1000 ? `${n}` : `${(n / 1000).toFixed(1)}k`;
					const usage = ctx.getContextUsage();
					const pct = usage?.percent ?? 0;
					const filled = Math.round(pct / 10) || 1;

					const l1Left =
						theme.fg("dim", " ") +
						theme.fg("accent", `${tier.icon} ${tier.name}`) +
						theme.fg("dim", ` [${tier.thinking}]${lockIcon} `) +
						theme.fg("warning", "[") +
						theme.fg("success", "#".repeat(filled)) +
						theme.fg("dim", "-".repeat(10 - filled)) +
						theme.fg("warning", "]") +
						theme.fg("dim", " ") +
						theme.fg("accent", `${Math.round(pct)}%`);

					const l1Right =
						theme.fg("success", `${fmt(tokIn)}`) +
						theme.fg("dim", " in ") +
						theme.fg("accent", `${fmt(tokOut)}`) +
						theme.fg("dim", " out ") +
						theme.fg("warning", `$${cost.toFixed(4)}`) +
						theme.fg("dim", " ");

					const pad1 = " ".repeat(Math.max(1, width - visibleWidth(l1Left) - visibleWidth(l1Right)));
					const line1 = truncateToWidth(l1Left + pad1 + l1Right, width, "");

					const dir = basename(ctx.cwd);
					const branch = footerData.getGitBranch();
					const lastReason = state.tierHistory.length > 0
						? state.tierHistory[state.tierHistory.length - 1].reason
						: "start: balanced";

					const l2Left =
						theme.fg("dim", ` ${dir}`) +
						(branch ? theme.fg("dim", " ") + theme.fg("warning", "(") + theme.fg("success", branch) + theme.fg("warning", ")") : "") +
						theme.fg("dim", " · ") +
						theme.fg("accent", `${state.totalSwitches}`) +
						theme.fg("dim", " switches") +
						(state.savedVsAlwaysOpus > 0 ? theme.fg("dim", " · saved ") + theme.fg("success", `$${state.savedVsAlwaysOpus.toFixed(4)}`) : "");

					const l2Right = theme.fg("dim", truncateToWidth(lastReason, Math.floor(width * 0.4), "…") + " ");

					const pad2 = " ".repeat(Math.max(1, width - visibleWidth(l2Left) - visibleWidth(l2Right)));
					const line2 = truncateToWidth(l2Left + pad2 + l2Right, width, "");

					return [line1, line2];
				},
			};
		});

		ctx.ui.notify(
			`Model Router active — LLM-powered (Haiku classifier)\n` +
			`Start: ${TIERS[1].icon} ${TIERS[1].name} · Auto-routes per prompt\n` +
			`/router — Status & controls`,
			"info",
		);
	});

	// ── /router Command ────────────────────────────────────────────────────

	pi.registerCommand("router", {
		description: "Model router status & controls. Usage: /router [lock|unlock|tier 1|2|3]",
		async handler(args: string, ctx) {
			const parts = args.trim().split(/\s+/);
			const sub = parts[0]?.toLowerCase();

			if (sub === "lock") { state.locked = true; ctx.ui.notify("🔒 Locked", "info"); return; }
			if (sub === "unlock") { state.locked = false; ctx.ui.notify("🔓 Unlocked", "info"); return; }
			if (sub === "tier") {
				const n = parseInt(parts[1], 10);
				if (n >= 1 && n <= 3) {
					await switchToTier(n - 1, `manual /router tier ${n}`, ctx);
					ctx.ui.notify(`${TIERS[n - 1].icon} ${TIERS[n - 1].name}`, "info");
					return;
				}
				ctx.ui.notify("Usage: /router tier 1|2|3", "warning"); return;
			}

			const tier = TIERS[state.currentTier];
			const history = state.tierHistory.length > 0
				? state.tierHistory.slice(-5).map(h => `    Turn ${h.turn}: ${h.reason}`).join("\n")
				: "    (none)";

			pi.sendMessage({
				customType: "router-status",
				content: [
					``,
					`  ┌─── Model Router (LLM-powered) ───┐`,
					`  │  Current: ${tier.icon} ${tier.name.padEnd(22)}│`,
					`  │  Thinking: ${tier.thinking.padEnd(21)}│`,
					`  │  Locked: ${(state.locked ? "yes 🔒" : "no").padEnd(23)}│`,
					`  │  Switches: ${String(state.totalSwitches).padEnd(21)}│`,
					`  │  Turns: ${String(state.turnCount).padEnd(23)}│`,
					`  │  Tool calls: ${String(state.totalToolCalls).padEnd(18)}│`,
					`  │  Errors: ${String(state.totalErrors).padEnd(22)}│`,
					`  │  Saved vs Opus: $${state.savedVsAlwaysOpus.toFixed(4).padEnd(14)}│`,
					`  │  Classifier cost: $${state.classifierCost.toFixed(4).padEnd(12)}│`,
					`  └────────────────────────────────────┘`,
					``,
					`  Recent routing:`,
					history,
					``,
					`  /router lock | unlock | tier 1|2|3`,
					``,
				].join("\n"),
				display: true,
			});
		},
	});
}
