/**
 * Nexus — The unified Pi intelligence layer
 *
 * Combines model routing + sub-agent orchestration into one seamless system.
 * The main agent stays fast and interactive. Complex work is automatically
 * offloaded to sub-agents. The model auto-adapts to task complexity.
 *
 * Model routing uses a REAL LLM call (Haiku, ~200ms) to classify task
 * complexity — no brittle regex patterns. The classifier returns a tier
 * (1-3) and whether to offload to a sub-agent.
 *
 * Features:
 *   ✦ LLM-powered task classification (Haiku as router brain)
 *   ✦ Auto model routing (Haiku → Sonnet → Opus based on classification)
 *   ✦ Sub-agent spawning for heavy work (preserves main context)
 *   ✦ Auto-delegation when classifier says to offload
 *   ✦ Context pressure monitoring (warns & suggests offload at 70%+)
 *   ✦ Mid-turn escalation on errors/heavy tool usage
 *   ✦ Live dashboard widget showing all agents
 *   ✦ Cost tracking across main + sub-agents
 *   ✦ Persistent sub-agent sessions for multi-turn continuations
 *
 * Usage:
 *   pi -e extensions/nexus.ts
 *   alias pn='pi -e extensions/nexus.ts'
 *
 * Commands:
 *   /nx              — Full dashboard
 *   /nx spawn <task> — Manual sub-agent spawn
 *   /nx kill <id>    — Kill sub-agent
 *   /nx clear        — Clear all sub-agents
 *   /nx mode <m>     — interactive | hybrid | delegator
 *   /nx auto on|off  — Toggle auto-delegation
 *   /nx lock         — Lock model (disable auto-routing)
 *   /nx unlock       — Unlock model
 *   /nx tier 1|2|3   — Force model tier
 */

import type { AssistantMessage as AssistantMsg } from "@mariozechner/pi-ai";
import { complete, getModel } from "@mariozechner/pi-ai";
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
import { DynamicBorder } from "@mariozechner/pi-coding-agent";
import { Container, Text, truncateToWidth, visibleWidth, type AutocompleteItem } from "@mariozechner/pi-tui";
import { Type } from "@sinclair/typebox";
import { spawn as cpSpawn, type ChildProcess } from "child_process";
import { existsSync, mkdirSync, readdirSync, readFileSync } from "fs";
import { basename, join } from "path";
import { applyExtensionDefaults } from "./themeMap.ts";

// ══════════════════════════════════════════════════════════════════════════
// CONSTANTS
// ══════════════════════════════════════════════════════════════════════════

const AGENT_TIMEOUT_MS = 10 * 60 * 1000;
const WIDGET_THROTTLE_MS = 400;
const CONTEXT_WARN_PCT = 70;
const MAX_RESULT_LEN = 12000;

// ══════════════════════════════════════════════════════════════════════════
// MODEL TIERS
// ══════════════════════════════════════════════════════════════════════════

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

const TIERS: Tier[] = [
	{ 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-POWERED TASK CLASSIFIER
// ══════════════════════════════════════════════════════════════════════════

const CLASSIFIER_SYSTEM_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>, "thinking": "<off|low|medium|high>", "offload": <true|false>, "reason": "<brief reason>"}

Classification rules:

TIER 1 (simple — use cheapest model):
- Quick questions, yes/no answers, explanations
- Reading/viewing files, checking status, listing things
- Small single-line edits, formatting, renaming
- Conversational responses (thanks, ok, etc.)
- Simple lookups or searches

TIER 2 (medium — use balanced model):
- Writing new functions, classes, or components
- Debugging, fixing bugs, error resolution
- Refactoring single files
- Writing tests for specific code
- Database queries, API endpoint creation
- Docker/CI configuration changes

TIER 3 (complex — use most powerful model):
- Architecture design, system planning
- Multi-file refactoring or migrations
- Security audits, performance optimization
- Complex algorithmic problems
- Designing new systems from scratch
- Tasks requiring deep reasoning or trade-off analysis

THINKING LEVEL (extended reasoning budget — set independently of tier):
- "off"    → Simple/conversational. No benefit from reasoning.
- "low"    → Moderate task. Light reasoning helpful (most Tier 2).
- "medium" → Tricky bug, subtle logic, multi-step plan (upper Tier 2 / lower Tier 3).
- "high"   → Architecture, security audit, complex algorithm, deep trade-offs (Tier 3).

Rule of thumb: Tier 1 → "off". Tier 2 → "low" or "medium". Tier 3 → "medium" or "high".
Only use "high" when extended chain-of-thought would meaningfully improve correctness.

OFFLOAD (spawn sub-agent instead of inline):
- Set true when the task would generate lots of context (multi-file changes, codebase-wide operations)
- Set true when the task is independent and can run in background
- Set true when "across all files", "entire codebase", parallel work is mentioned
- Set false for conversational, quick, or interactive tasks
- Set false when user needs immediate back-and-forth

Be concise. The "reason" should be under 10 words.`;

interface ClassificationResult {
	tier: number;
	thinking: "off" | "low" | "medium" | "high";
	offload: boolean;
	reason: string;
}

// Cache recent classifications to avoid duplicate calls
const classificationCache = new Map<string, { result: ClassificationResult; timestamp: number }>();
const CACHE_TTL_MS = 60_000; // 1 minute

async function classifyWithLLM(
	prompt: string,
	ctx: ExtensionContext,
	contextPercent: number | null,
): Promise<ClassificationResult> {
	// Check cache (normalize by trimming + lowercasing)
	const cacheKey = prompt.trim().toLowerCase().slice(0, 200);
	const cached = classificationCache.get(cacheKey);
	if (cached && Date.now() - cached.timestamp < CACHE_TTL_MS) {
		return cached.result;
	}

	// Use Haiku for classification — fast and cheap (~0.001 cents per classification)
	const classifier = ctx.modelRegistry.find("anthropic", "claude-haiku-4-5");
	if (!classifier) {
		// Fallback: simple heuristic if Haiku unavailable
		return { tier: 2, thinking: "low", offload: false, reason: "classifier unavailable" };
	}

	const apiKey = await ctx.modelRegistry.getApiKey(classifier);
	if (!apiKey) {
		return { tier: 2, thinking: "low", offload: false, reason: "no API key" };
	}

	// Build context message with additional signals
	let userMsg = prompt;
	if (contextPercent !== null && contextPercent > CONTEXT_WARN_PCT) {
		userMsg += `\n\n[SYSTEM NOTE: Main context is at ${Math.round(contextPercent)}% capacity. Prefer offload=true for heavy tasks.]`;
	}

	try {
		const response = await complete(classifier, {
			systemPrompt: CLASSIFIER_SYSTEM_PROMPT,
			messages: [{
				role: "user" as const,
				content: userMsg,
				timestamp: Date.now(),
			}],
		}, {
			reasoning: "off",
		});

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

		// Parse JSON from response (handle potential markdown wrapping)
		const jsonStr = text.replace(/```json?\n?/g, "").replace(/```/g, "").trim();
		const parsed = JSON.parse(jsonStr);

		const validThinking = ["off", "low", "medium", "high"];
		const result: ClassificationResult = {
			tier: Math.max(1, Math.min(3, parsed.tier || 2)),
			thinking: validThinking.includes(parsed.thinking) ? parsed.thinking : "low",
			offload: Boolean(parsed.offload),
			reason: String(parsed.reason || "classified").slice(0, 50),
		};

		// Cache it
		classificationCache.set(cacheKey, { result, timestamp: Date.now() });

		return result;
	} catch (err: any) {
		// On any failure, default to Sonnet, no offload
		return { tier: 2, thinking: "low", offload: false, reason: `classify error: ${err?.message?.slice(0, 30)}` };
	}
}

// ══════════════════════════════════════════════════════════════════════════
// SUB-AGENT TYPES
// ══════════════════════════════════════════════════════════════════════════

type AgentMode = "interactive" | "delegator" | "hybrid";
type SubStatus = "queued" | "running" | "done" | "error" | "killed";

interface SubAgent {
	id: number;
	status: SubStatus;
	task: string;
	role: string;
	tools: string;
	textChunks: string[];
	toolCount: number;
	elapsed: number;
	startTime: number;
	sessionFile: string;
	turnCount: number;
	proc?: ChildProcess;
	timer?: ReturnType<typeof setInterval>;
	inputTokens: number;
	outputTokens: number;
	cost: number;
}

interface RoleDef {
	name: string;
	description: string;
	tools: string;
	systemPrompt: string;
}

const ROLES: RoleDef[] = [
	{ name: "general", description: "General-purpose coding agent", tools: "read,bash,edit,write,grep,find,ls", systemPrompt: "You are a focused coding agent. Complete the given task efficiently. Be thorough but concise." },
	{ name: "scout", description: "Fast recon (read-only)", tools: "read,grep,find,ls", systemPrompt: "You are a scout. Investigate the codebase quickly and report findings. Do NOT modify files." },
	{ name: "builder", description: "Implementation & code generation", tools: "read,write,edit,bash,grep,find,ls", systemPrompt: "You are a builder. Implement changes thoroughly. Clean, minimal code. Follow existing patterns." },
	{ name: "reviewer", description: "Code review (read-only)", tools: "read,grep,find,ls", systemPrompt: "You are a reviewer. Analyze code for bugs, security, performance, style. Be specific about line numbers." },
	{ name: "tester", description: "Test writing & execution", tools: "read,write,edit,bash,grep,find,ls", systemPrompt: "You are a tester. Write and run tests. Cover edge cases. Use the project's existing test framework." },
];

// ══════════════════════════════════════════════════════════════════════════
// EXTENSION
// ══════════════════════════════════════════════════════════════════════════

export default function nexus(pi: ExtensionAPI) {

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

	let currentTier = 1;
	let lastThinking: "off" | "low" | "medium" | "high" = "low";
	let tierLocked = false;
	let totalSwitches = 0;
	let turnsSinceSwitch = 0;
	let consecutiveErrors = 0;
	let totalToolCalls = 0;
	let editWriteThisTurn = 0;
	let toolsThisTurn = 0;
	let turnCount = 0;
	let savedVsOpus = 0;
	let classifierCost = 0; // Track cost of Haiku classifier calls
	const tierHistory: Array<{ tier: number; reason: string; turn: number }> = [];

	// ── Agent State ────────────────────────────────────────────────────────

	const agents: Map<number, SubAgent> = new Map();
	const activeProcs: Set<ChildProcess> = new Set();
	let nextId = 1;
	let widgetCtx: ExtensionContext | null = null;
	let mode: AgentMode = "hybrid";
	let autoDelegate = true;
	let sessionDir = "";
	let customRoles: RoleDef[] = [];

	// ── Widget throttle ────────────────────────────────────────────────────

	let wDirty = false;
	let wTimer: ReturnType<typeof setTimeout> | null = null;

	function scheduleWidget() {
		wDirty = true;
		if (wTimer) return;
		wTimer = setTimeout(() => { wTimer = null; if (wDirty) { wDirty = false; renderWidget(); } }, WIDGET_THROTTLE_MS);
	}

	function flushWidget() {
		if (wTimer) { clearTimeout(wTimer); wTimer = null; }
		wDirty = false;
		renderWidget();
	}

	// ── Helpers ────────────────────────────────────────────────────────────

	function getAllRoles(): RoleDef[] { return [...ROLES, ...customRoles]; }

	function findRole(name: string): RoleDef {
		return getAllRoles().find(r => r.name.toLowerCase() === name.toLowerCase()) || ROLES[0];
	}

	function makeSessionFile(id: number): string {
		if (!existsSync(sessionDir)) mkdirSync(sessionDir, { recursive: true });
		return join(sessionDir, `nexus-${id}-${Date.now()}.jsonl`);
	}

	function loadCustomRoles(cwd: string) {
		customRoles = [];
		const seen = new Set(ROLES.map(r => r.name));
		for (const dir of [join(cwd, ".pi", "agents"), join(cwd, "agents")]) {
			if (!existsSync(dir)) continue;
			try {
				for (const file of readdirSync(dir)) {
					if (!file.endsWith(".md")) continue;
					const raw = readFileSync(join(dir, file), "utf-8");
					const m = raw.match(/^---\n([\s\S]*?)\n---\n([\s\S]*)$/);
					if (!m) continue;
					const fm: Record<string, string> = {};
					for (const line of m[1].split("\n")) {
						const i = line.indexOf(":");
						if (i > 0) fm[line.slice(0, i).trim()] = line.slice(i + 1).trim();
					}
					if (!fm.name || seen.has(fm.name)) continue;
					seen.add(fm.name);
					customRoles.push({ name: fm.name, description: fm.description || "", tools: fm.tools || "read,grep,find,ls", systemPrompt: m[2].trim() });
				}
			} catch {}
		}
	}

	// ── Model Switching ────────────────────────────────────────────────────

	async function switchTier(target: number, reason: string, ctx?: ExtensionContext, thinkingOverride?: "off" | "low" | "medium" | "high") {
		const t = Math.max(0, Math.min(2, target));
		const tier = TIERS[t];
		const model = ctx?.modelRegistry.find(tier.provider, tier.modelId);
		if (!model) return;

		const thinking = thinkingOverride ?? tier.thinking;
		const tierChanged = t !== currentTier;
		const thinkingChanged = thinking !== lastThinking;

		if (!tierChanged && !thinkingChanged) return;

		if (tierChanged) {
			const ok = await pi.setModel(model);
			if (!ok) return;
		}
		pi.setThinkingLevel(thinking);
		lastThinking = thinking;

		const prev = currentTier;
		currentTier = t;
		if (tierChanged) totalSwitches++;
		turnsSinceSwitch = 0;
		const thinkTag = thinkingOverride ? ` 🧠${thinking}` : "";
		tierHistory.push({ tier: t, reason: `${TIERS[prev].icon}→${tier.icon} ${reason}${thinkTag}`, turn: turnCount });
	}

	// ── Widget Rendering ───────────────────────────────────────────────────

	function renderWidget() {
		if (!widgetCtx) return;
		if (agents.size === 0) {
			widgetCtx.ui.setWidget("nexus-agents", undefined);
			return;
		}

		widgetCtx.ui.setWidget("nexus-agents", (_tui: any, theme: any) => {
			const container = new Container();
			const content = new Text("", 0, 0);
			container.addChild(content);

			return {
				render(width: number): string[] {
					const lines: string[] = [];
					const tier = TIERS[currentTier];
					const modeIcon = mode === "hybrid" ? "⚡" : mode === "delegator" ? "📡" : "💬";

					lines.push(theme.fg("dim", "─".repeat(width)));
					lines.push(
						theme.fg("accent", ` ✦ Nexus`) +
						theme.fg("dim", ` ${tier.icon} ${tier.name} · ${modeIcon} ${mode} · `) +
						theme.fg("accent", `${agents.size}`) +
						theme.fg("dim", " agents")
					);

					for (const [, a] of agents) {
						const icon = a.status === "running" ? "●" : a.status === "done" ? "✓" : a.status === "queued" ? "◌" : "✗";
						const color = a.status === "running" ? "accent" : a.status === "done" ? "success" : a.status === "queued" ? "dim" : "error";
						const elapsed = a.status === "running" ? Math.round((Date.now() - a.startTime) / 1000) : Math.round(a.elapsed / 1000);
						const task = a.task.length > 50 ? a.task.slice(0, 47) + "..." : a.task;
						const turn = a.turnCount > 1 ? ` t${a.turnCount}` : "";

						lines.push(
							theme.fg(color, `  ${icon} #${a.id}`) +
							theme.fg("dim", ` [${a.role}${turn}] ${elapsed}s · ${a.toolCount} tools`) +
							(a.cost > 0 ? theme.fg("warning", ` $${a.cost.toFixed(4)}`) : "")
						);
						lines.push(theme.fg("muted", `    ${task}`));

						if (a.status === "running" && a.textChunks.length > 0) {
							const last = a.textChunks.join("").split("\n").filter((l: string) => l.trim()).pop() || "";
							if (last) lines.push(theme.fg("dim", `    → ${last.slice(0, width - 8)}`));
						}
					}

					lines.push(theme.fg("dim", "─".repeat(width)));
					content.setText(lines.join("\n"));
					return container.render(width);
				},
				invalidate() { container.invalidate(); },
			};
		});
	}

	// ── Sub-Agent Process ──────────────────────────────────────────────────

	function processLine(a: SubAgent, line: string) {
		if (!line.trim()) return;
		try {
			const e = JSON.parse(line);
			if (e.type === "message_update" && e.assistantMessageEvent?.type === "text_delta") {
				a.textChunks.push(e.assistantMessageEvent.delta || "");
				scheduleWidget();
			} else if (e.type === "tool_execution_start") {
				a.toolCount++;
				scheduleWidget();
			} else if (e.type === "message_end" || e.type === "agent_end") {
				const msg = e.message || (e.messages || []).reverse().find((m: any) => m.role === "assistant");
				if (msg?.usage) {
					a.inputTokens = msg.usage.input || 0;
					a.outputTokens = msg.usage.output || 0;
					a.cost = msg.usage.cost?.total || 0;
					scheduleWidget();
				}
			}
		} catch {}
	}

	function spawnProcess(a: SubAgent, prompt: string, ctx: ExtensionContext, isCont: boolean): Promise<{ output: string; exitCode: number }> {
		a.status = "running";
		a.startTime = Date.now();
		scheduleWidget();

		const role = findRole(a.role);
		// Sub-agents use Sonnet for balanced cost/performance
		const model = "anthropic/claude-sonnet-4-6";

		const args = [
			"--mode", "json", "-p", "--no-extensions", "--no-skills",
			"--model", model, "--tools", role.tools, "--thinking", "low",
			"--append-system-prompt", role.systemPrompt,
			"--session", a.sessionFile,
		];
		if (isCont) args.push("-c");
		args.push(prompt);

		return new Promise(resolve => {
			let done = false;
			const fin = (val: { output: string; exitCode: number }) => { if (!done) { done = true; resolve(val); } };

			const proc = cpSpawn("pi", args, { stdio: ["ignore", "pipe", "pipe"], env: { ...process.env } });
			a.proc = proc;
			activeProcs.add(proc);

			a.timer = setInterval(() => { a.elapsed = Date.now() - a.startTime; scheduleWidget(); }, 1000);

			const timeout = setTimeout(() => {
				try { proc.kill("SIGTERM"); } catch {}
				setTimeout(() => { try { proc.kill("SIGKILL"); } catch {} }, 3000);
			}, AGENT_TIMEOUT_MS);

			let buf = "";
			proc.stdout!.setEncoding("utf-8");
			proc.stdout!.on("data", (chunk: string) => {
				buf += chunk;
				const lines = buf.split("\n");
				buf = lines.pop() || "";
				for (const l of lines) processLine(a, l);
			});
			proc.stderr!.setEncoding("utf-8");
			proc.stderr!.on("data", () => {});

			proc.on("close", code => {
				clearTimeout(timeout);
				activeProcs.delete(proc);
				a.proc = undefined;
				if (a.timer) clearInterval(a.timer);
				a.elapsed = Date.now() - a.startTime;
				if (buf.trim()) processLine(a, buf);
				a.status = a.elapsed >= AGENT_TIMEOUT_MS ? "error" : (code === 0 ? "done" : "error");
				flushWidget();
				fin({ output: a.textChunks.join(""), exitCode: code ?? 1 });
			});

			proc.on("error", err => {
				clearTimeout(timeout);
				activeProcs.delete(proc);
				a.proc = undefined;
				if (a.timer) clearInterval(a.timer);
				a.status = "error";
				flushWidget();
				fin({ output: `Error: ${err.message}`, exitCode: 1 });
			});
		});
	}

	function spawnAndDeliver(a: SubAgent, prompt: string, ctx: ExtensionContext, isCont: boolean) {
		spawnProcess(a, prompt, ctx, isCont).then(({ output, exitCode }) => {
			const status = exitCode === 0 ? "✓" : "✗";
			const elapsed = Math.round(a.elapsed / 1000);
			const turn = a.turnCount > 1 ? ` (Turn ${a.turnCount})` : "";
			const truncated = output.length > MAX_RESULT_LEN ? output.slice(0, MAX_RESULT_LEN) + "\n\n... [truncated]" : output;

			ctx.ui.notify(`#${a.id} [${a.role}] ${status} ${elapsed}s`, exitCode === 0 ? "info" : "error");

			pi.sendMessage({
				customType: "nexus-result",
				content: `**Sub-agent #${a.id}** [${a.role}]${turn} — ${status} ${elapsed}s | ${a.toolCount} tools | $${a.cost.toFixed(4)}\n\n${truncated}`,
				display: "assistant",
			}, { deliverAs: "followUp", triggerTurn: true });
		});
	}

	// ── LLM Tools ──────────────────────────────────────────────────────────

	const roleList = ROLES.map(r => `${r.name} (${r.description})`).join(", ");

	pi.registerTool({
		name: "nexus_spawn",
		label: "Spawn Sub-Agent",
		description: `Spawn a background sub-agent for heavy tasks. Preserves main chat context. Results auto-deliver when done.\n\nRoles: ${roleList}. Custom roles from .pi/agents/*.md also available.`,
		promptGuidelines: [
			"Spawn sub-agents for multi-file, long-running, or context-heavy work",
			"Use role='scout' for exploration, 'builder' for code, 'reviewer' for review, 'tester' for tests",
			"Spawn multiple agents in parallel for independent tasks",
			"Continue chatting while sub-agents work — results arrive automatically",
		],
		parameters: Type.Object({
			task: Type.String({ description: "Complete task description — sub-agent has no context from this conversation" }),
			role: Type.Optional(Type.String({ description: "Agent role. Default: general" })),
		}),
		async execute(_id, params, _sig, onUpdate, ctx) {
			widgetCtx = ctx;
			const id = nextId++;
			const a: SubAgent = {
				id, status: "queued", task: params.task, role: params.role || "general",
				tools: findRole(params.role || "general").tools, textChunks: [], toolCount: 0,
				elapsed: 0, startTime: Date.now(), sessionFile: makeSessionFile(id),
				turnCount: 1, inputTokens: 0, outputTokens: 0, cost: 0,
			};
			agents.set(id, a);
			if (onUpdate) onUpdate({ content: [{ type: "text", text: `Spawning #${id} [${a.role}]...` }], details: { id, role: a.role } });
			spawnAndDeliver(a, params.task, ctx, false);
			return { content: [{ type: "text", text: `Sub-agent #${id} [${a.role}] running in background. Continue chatting.` }], details: { id, role: a.role } };
		},
		renderCall(args, theme) {
			const r = (args as any).role || "general";
			const t = ((args as any).task || "").slice(0, 55);
			return new Text(theme.fg("toolTitle", theme.bold("nexus_spawn ")) + theme.fg("accent", `[${r}] `) + theme.fg("muted", t), 0, 0);
		},
		renderResult(result, opts, theme) {
			const d = result.details as any;
			if (!d) return undefined;
			return new Text(theme.fg("success", `◉ #${d.id}`) + theme.fg("dim", ` [${d.role}] background`), 0, 0);
		},
	});

	pi.registerTool({
		name: "nexus_continue",
		label: "Continue Sub-Agent",
		description: "Continue a finished sub-agent's conversation with follow-up instructions.",
		parameters: Type.Object({
			id: Type.Number({ description: "Sub-agent ID" }),
			prompt: Type.String({ description: "Follow-up instructions" }),
		}),
		async execute(_cid, params, _sig, _upd, ctx) {
			widgetCtx = ctx;
			const a = agents.get(params.id);
			if (!a) return { content: [{ type: "text", text: `No #${params.id} found.` }] };
			if (a.status === "running") return { content: [{ type: "text", text: `#${params.id} still running.` }] };
			a.textChunks = []; a.toolCount = 0; a.elapsed = 0; a.turnCount++;
			spawnAndDeliver(a, params.prompt, ctx, true);
			return { content: [{ type: "text", text: `#${params.id} continuing (Turn ${a.turnCount}).` }] };
		},
	});

	pi.registerTool({
		name: "nexus_status",
		label: "Agent Status",
		description: "Check all sub-agents.",
		parameters: Type.Object({}),
		async execute() {
			if (agents.size === 0) return { content: [{ type: "text", text: "No sub-agents." }] };
			const lines = Array.from(agents.values()).map(a => {
				const e = a.status === "running" ? Math.round((Date.now() - a.startTime) / 1000) : Math.round(a.elapsed / 1000);
				return `#${a.id} [${a.role}] ${a.status} — ${e}s, ${a.toolCount} tools, $${a.cost.toFixed(4)}\n  ${a.task}`;
			});
			return { content: [{ type: "text", text: lines.join("\n") }] };
		},
	});

	pi.registerTool({
		name: "nexus_kill",
		label: "Kill Sub-Agent",
		description: "Kill/remove a sub-agent.",
		parameters: Type.Object({ id: Type.Number() }),
		async execute(_cid, params) {
			const a = agents.get(params.id);
			if (!a) return { content: [{ type: "text", text: `No #${params.id}.` }] };
			if (a.proc) try { a.proc.kill("SIGTERM"); } catch {}
			if (a.timer) clearInterval(a.timer);
			agents.delete(params.id);
			flushWidget();
			return { content: [{ type: "text", text: `#${params.id} removed.` }] };
		},
	});

	pi.registerTool({
		name: "nexus_collect",
		label: "Collect Results",
		description: "Collect results from all finished sub-agents.",
		parameters: Type.Object({}),
		async execute() {
			const done = Array.from(agents.values()).filter(a => a.status === "done" || a.status === "error");
			if (done.length === 0) return { content: [{ type: "text", text: "No results to collect." }] };
			let totalCost = 0;
			const parts = done.map(a => {
				const out = a.textChunks.join("");
				const trunc = out.length > 6000 ? out.slice(0, 6000) + "\n... [truncated]" : out;
				totalCost += a.cost;
				return `## ${a.status === "done" ? "✓" : "✗"} #${a.id} [${a.role}] ${Math.round(a.elapsed / 1000)}s\n${a.task}\n\n${trunc}`;
			});
			return { content: [{ type: "text", text: `${done.length} results ($${totalCost.toFixed(4)})\n\n${parts.join("\n\n---\n\n")}` }] };
		},
	});

	// ══════════════════════════════════════════════════════════════════════
	// CORE ROUTING LOGIC — LLM-POWERED
	// ══════════════════════════════════════════════════════════════════════

	pi.on("before_agent_start", async (event, ctx) => {
		widgetCtx = ctx;
		const usage = ctx.getContextUsage();
		const contextPct = usage?.percent ?? null;

		// ── Step 1: LLM Classification ─────────────────────────────────────
		let classification: ClassificationResult | null = null;

		if (!tierLocked || (autoDelegate && mode !== "interactive")) {
			ctx.ui.setWorkingMessage("classifying task...");
			classification = await classifyWithLLM(event.prompt, ctx, contextPct);
			ctx.ui.setWorkingMessage(); // restore default

			// Track classifier cost (Haiku: ~150 input + ~30 output tokens per call)
			classifierCost += (150 * 0.80 + 30 * 4.00) / 1_000_000;
		}

		// ── Step 2: Model Routing ──────────────────────────────────────────
		if (!tierLocked && classification) {
			let target = classification.tier - 1; // Convert 1-3 to 0-2
			let thinking = classification.thinking;

			// Context-aware adjustments
			// If we've been doing heavy tool work, don't drop to Haiku
			if (target === 0 && totalToolCalls > 10) target = 1;

			// Error escalation override
			if (consecutiveErrors >= 2 && target < 2) {
				target = Math.min(target + 1, 2);
				if (thinking === "off" || thinking === "low") thinking = "medium";
			}

			// Context pressure: bump tier if context is getting full
			if (contextPct !== null && contextPct > 80 && target < 1) {
				target = 1; // At least Sonnet when context is pressured
			}

			await switchTier(target, classification.reason, ctx, thinking);
		}

		// ── Step 3: Auto-Delegation ────────────────────────────────────────
		if (classification?.offload && autoDelegate && (mode === "hybrid" || mode === "delegator")) {
			// Context pressure warning
			if (contextPct !== null && contextPct > CONTEXT_WARN_PCT) {
				ctx.ui.setStatus("nexus-ctx", `⚠️ Context ${Math.round(contextPct)}%`);
			} else {
				ctx.ui.setStatus("nexus-ctx", undefined);
			}

			const roles = getAllRoles().map(r => `${r.name}: ${r.description}`).join("\n");
			const contextWarn = (contextPct !== null && contextPct > CONTEXT_WARN_PCT)
				? `\n⚠️ CONTEXT AT ${Math.round(contextPct)}% — strongly prefer sub-agents for this work.\n`
				: "";

			return {
				systemPrompt: event.systemPrompt + `\n\n## Nexus: Sub-Agent Delegation Recommended
The task classifier determined this task should be **offloaded to a sub-agent** (reason: "${classification.reason}").

You have these tools: nexus_spawn, nexus_continue, nexus_status, nexus_kill, nexus_collect.
${contextWarn}
**Strongly prefer spawning a sub-agent** for this task rather than doing it inline. This preserves the main chat context for continued interaction.

Available roles:\n${roles}

Sub-agents run in background with their own context. You can spawn multiple for parallel work. Results auto-deliver. Keep your main response brief — explain what you're delegating and why.`,
			};
		} else {
			ctx.ui.setStatus("nexus-ctx", undefined);
		}
	});

	// ── Turn & Tool Tracking (mid-conversation escalation) ─────────────────

	pi.on("turn_start", async () => {
		toolsThisTurn = 0;
		editWriteThisTurn = 0;
		turnCount++;
		turnsSinceSwitch++;
	});

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

	pi.on("tool_execution_end", async (event, ctx) => {
		if (event.isError) {
			consecutiveErrors++;
			// Mid-turn escalation on repeated errors
			if (consecutiveErrors >= 2 && !tierLocked && currentTier < 2) {
				await switchTier(currentTier + 1, `${consecutiveErrors} consecutive errors`, ctx);
			}
		} else {
			consecutiveErrors = 0;
		}
	});

	pi.on("turn_end", async (_event, ctx) => {
		if (!tierLocked) {
			// Escalate on heavy tool usage
			if (editWriteThisTurn >= 4 && currentTier < 2) {
				await switchTier(2, `heavy: ${editWriteThisTurn} edits/writes`, ctx);
			} else if (toolsThisTurn >= 6 && currentTier < 1) {
				await switchTier(1, `busy: ${toolsThisTurn} tool calls`, ctx);
			}

			// De-escalate after quiet period (but only if classifier agrees next time)
			// We just track turns since switch — the classifier handles de-escalation naturally
		}

		// Cost savings estimate
		const opusCost = (2000 * TIERS[2].inputCost + 1000 * TIERS[2].outputCost) / 1e6;
		const actualCost = (2000 * TIERS[currentTier].inputCost + 1000 * TIERS[currentTier].outputCost) / 1e6;
		savedVsOpus += opusCost - actualCost;
	});

	// ── /nx Command ────────────────────────────────────────────────────────

	pi.registerCommand("nx", {
		description: "Nexus controls: /nx [spawn|kill|clear|mode|auto|lock|unlock|tier]",
		getArgumentCompletions: (prefix: string): AutocompleteItem[] | null => {
			const items = [
				{ value: "spawn ", label: "spawn <task>" },
				{ value: "kill ", label: "kill <id>" },
				{ value: "clear", label: "clear all agents" },
				{ value: "mode hybrid", label: "mode hybrid" },
				{ value: "mode interactive", label: "mode interactive" },
				{ value: "mode delegator", label: "mode delegator" },
				{ value: "auto on", label: "auto on" },
				{ value: "auto off", label: "auto off" },
				{ value: "lock", label: "lock model" },
				{ value: "unlock", label: "unlock model" },
				{ value: "tier ", label: "tier 1|2|3" },
			];
			return items.filter(i => i.value.startsWith(prefix));
		},

		async handler(args, ctx) {
			widgetCtx = ctx;
			const parts = (args || "").trim().split(/\s+/);
			const sub = parts[0]?.toLowerCase();

			if (!sub) {
				const tier = TIERS[currentTier];
				const modeIcon = mode === "hybrid" ? "⚡" : mode === "delegator" ? "📡" : "💬";
				const usage = ctx.getContextUsage();
				const pct = usage?.percent ? Math.round(usage.percent) : "?";
				const agentLines = agents.size === 0 ? "  (none)" : Array.from(agents.values()).map(a => {
					const e = a.status === "running" ? Math.round((Date.now() - a.startTime) / 1000) : Math.round(a.elapsed / 1000);
					return `  #${a.id} [${a.role}] ${a.status} ${e}s $${a.cost.toFixed(4)} — ${a.task.slice(0, 50)}`;
				}).join("\n");

				const lastRoute = tierHistory.length > 0
					? tierHistory.slice(-3).map(h => `  Turn ${h.turn}: ${h.reason}`).join("\n")
					: "  (none yet)";

				pi.sendMessage({
					customType: "nexus-dashboard",
					content: [
						``,
						`✦ Nexus Dashboard`,
						`  Model: ${tier.icon} ${tier.name} [${lastThinking}]${tierLocked ? " 🔒" : ""}`,
						`  Mode: ${modeIcon} ${mode} | Auto: ${autoDelegate ? "on" : "off"} | Context: ${pct}%`,
						`  Switches: ${totalSwitches} | Saved: $${savedVsOpus.toFixed(4)} | Classifier: $${classifierCost.toFixed(4)}`,
						``,
						`  Sub-Agents (${agents.size}):`,
						agentLines,
						``,
						`  Recent Routing:`,
						lastRoute,
						``,
						`  /nx spawn|kill|clear|mode|auto|lock|unlock|tier`,
						``,
					].join("\n"),
					display: "assistant",
				});
				return;
			}

			if (sub === "spawn") {
				const task = parts.slice(1).join(" ").trim();
				if (!task) { ctx.ui.notify("Usage: /nx spawn <task>", "error"); return; }
				const id = nextId++;
				const a: SubAgent = { id, status: "queued", task, role: "general", tools: ROLES[0].tools, textChunks: [], toolCount: 0, elapsed: 0, startTime: Date.now(), sessionFile: makeSessionFile(id), turnCount: 1, inputTokens: 0, outputTokens: 0, cost: 0 };
				agents.set(id, a);
				spawnAndDeliver(a, task, ctx, false);
				ctx.ui.notify(`#${id} spawned`, "info");
			} else if (sub === "kill") {
				const id = parseInt(parts[1], 10);
				const a = agents.get(id);
				if (!a) { ctx.ui.notify(`No #${id}`, "error"); return; }
				if (a.proc) try { a.proc.kill("SIGTERM"); } catch {}
				if (a.timer) clearInterval(a.timer);
				agents.delete(id);
				flushWidget();
				ctx.ui.notify(`#${id} killed`, "info");
			} else if (sub === "clear") {
				for (const [, a] of agents) {
					if (a.proc) try { a.proc.kill("SIGTERM"); } catch {}
					if (a.timer) clearInterval(a.timer);
				}
				agents.clear();
				nextId = 1;
				flushWidget();
				ctx.ui.notify("Cleared", "info");
			} else if (sub === "mode") {
				const v = parts[1]?.toLowerCase() as AgentMode;
				if (["interactive", "hybrid", "delegator"].includes(v)) { mode = v; ctx.ui.notify(`Mode: ${v}`, "info"); }
				else ctx.ui.notify("Usage: /nx mode interactive|hybrid|delegator", "error");
			} else if (sub === "auto") {
				if (parts[1] === "on") { autoDelegate = true; ctx.ui.notify("Auto ON", "info"); }
				else if (parts[1] === "off") { autoDelegate = false; ctx.ui.notify("Auto OFF", "info"); }
				else ctx.ui.notify("Usage: /nx auto on|off", "error");
			} else if (sub === "lock") {
				tierLocked = true; ctx.ui.notify("🔒 Model locked", "info");
			} else if (sub === "unlock") {
				tierLocked = false; ctx.ui.notify("🔓 Model unlocked", "info");
			} else if (sub === "tier") {
				const n = parseInt(parts[1], 10);
				if (n >= 1 && n <= 3) {
					await switchTier(n - 1, `manual /nx tier ${n}`, ctx);
					ctx.ui.notify(`${TIERS[n - 1].icon} ${TIERS[n - 1].name}`, "info");
				} else ctx.ui.notify("Usage: /nx tier 1|2|3", "error");
			} else {
				ctx.ui.notify("Unknown. Try: /nx spawn|kill|clear|mode|auto|lock|unlock|tier", "error");
			}
		},
	});

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

	pi.on("session_start", async (_event, ctx) => {
		applyExtensionDefaults(import.meta.url, ctx);
		widgetCtx = ctx;
		sessionDir = join(ctx.cwd, ".pi", "agent-sessions", "nexus");
		loadCustomRoles(ctx.cwd);

		// Cleanup previous
		for (const [, a] of agents) {
			if (a.proc) try { a.proc.kill("SIGTERM"); } catch {}
			if (a.timer) clearInterval(a.timer);
		}
		agents.clear();
		nextId = 1;

		// Set initial model (Sonnet = balanced start)
		const init = TIERS[1];
		const model = ctx.modelRegistry.find(init.provider, init.modelId);
		if (model) {
			await pi.setModel(model);
			pi.setThinkingLevel(init.thinking);
		}
		currentTier = 1;

		// Footer
		ctx.ui.setFooter((_tui, theme, footerData) => {
			const unsub = footerData.onBranchChange(() => _tui.requestRender());
			return {
				dispose: unsub,
				invalidate() {},
				render(width: number): string[] {
					const tier = TIERS[currentTier];
					const lockStr = tierLocked ? " 🔒" : "";
					const usage = ctx.getContextUsage();
					const pct = usage?.percent ?? 0;
					const filled = Math.round(pct / 10) || 1;

					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 AssistantMsg;
							tokIn += m.usage.input;
							tokOut += m.usage.output;
							cost += m.usage.cost.total;
						}
					}
					// Add sub-agent costs + classifier cost
					for (const [, a] of agents) cost += a.cost;
					cost += classifierCost;
					const fmt = (n: number) => n < 1000 ? `${n}` : `${(n / 1000).toFixed(1)}k`;

					const l1L = theme.fg("dim", " ") + theme.fg("accent", `${tier.icon} ${tier.name}`) +
						theme.fg("dim", ` [${lastThinking}]${lockStr} `) +
						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 l1R = theme.fg("success", fmt(tokIn)) + theme.fg("dim", " in ") +
						theme.fg("accent", fmt(tokOut)) + theme.fg("dim", " out ") +
						theme.fg("warning", `$${cost.toFixed(4)} `);
					const p1 = " ".repeat(Math.max(1, width - visibleWidth(l1L) - visibleWidth(l1R)));

					const dir = basename(ctx.cwd);
					const branch = footerData.getGitBranch();
					const modeIcon = mode === "hybrid" ? "⚡" : mode === "delegator" ? "📡" : "💬";
					const running = Array.from(agents.values()).filter(a => a.status === "running").length;

					const l2L = theme.fg("dim", ` ${dir}`) +
						(branch ? theme.fg("dim", " ") + theme.fg("warning", "(") + theme.fg("success", branch) + theme.fg("warning", ")") : "") +
						theme.fg("dim", ` · ${modeIcon} ${mode}`);
					const l2R = (running > 0 ? theme.fg("accent", `● ${running} running `) : "") +
						(agents.size > 0 ? theme.fg("dim", `${agents.size} agents `) : "") +
						(savedVsOpus > 0 ? theme.fg("success", `saved $${savedVsOpus.toFixed(3)} `) : "");
					const p2 = " ".repeat(Math.max(1, width - visibleWidth(l2L) - visibleWidth(l2R)));

					return [
						truncateToWidth(l1L + p1 + l1R, width, ""),
						truncateToWidth(l2L + p2 + l2R, width, ""),
					];
				},
			};
		});

		const roles = getAllRoles().map(r => r.name).join(", ");
		ctx.ui.notify(
			`✦ Nexus active [${mode}] — LLM-powered routing\n` +
			`Model: ${TIERS[1].icon} ${TIERS[1].name} (Haiku classifies → auto-routes)\n` +
			`Roles: ${roles}\n` +
			`/nx — Dashboard & controls`,
			"info",
		);
	});

	// ── Cleanup ────────────────────────────────────────────────────────────

	function killAll() {
		for (const p of activeProcs) try { p.kill("SIGTERM"); } catch {}
		setTimeout(() => { for (const p of activeProcs) try { p.kill("SIGKILL"); } catch {} }, 3000);
	}
	pi.on("session_shutdown", async () => killAll());
	process.on("exit", killAll);
	process.on("SIGINT", () => { killAll(); process.exit(0); });
	process.on("SIGTERM", () => { killAll(); process.exit(0); });
}