/**
 * Minimal subagents extension.
 *
 * Registers a single `subagent` tool. Agent definitions are loaded from the
 * shared Pi user agents directory: ~/.pi/agent/agents.
 * Supports single and parallel execution. Output is verbal only (no file handoff).
 */
import { spawn } from "node:child_process";
import { fileURLToPath } from "node:url";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
import {
    getMarkdownTheme,
    truncateHead,
    truncateTail,
    withFileMutationQueue,
    DEFAULT_MAX_BYTES,
    DEFAULT_MAX_LINES,
} from "@earendil-works/pi-coding-agent";
import { Container, Markdown, Spacer, Text, truncateToWidth, visibleWidth, wrapTextWithAnsi, type Component } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { agentPath } from "../shared/paths.ts";
import { loadAgentDefinitions } from "../shared/agents.ts";

// ── Types ──────────────────────────────────────────────────────────────────────

export interface AgentConfig {
    name: string;
    description: string;
    tools: string[];
    model: string;
    thinking: string;
    systemPrompt: string;
    filePath: string;
    role?: "foreground" | "background" | "both";
    /**
     * If this agent has the `subagent` tool, restrict which agents it may spawn.
     * Passed to the child pi process via `PI_SUBAGENT_ALLOWED` so the child's
     * subagents extension filters its own registry before exposing it to the LLM.
     * `undefined` means no restriction (child sees every registered agent).
     */
    subagentAgents?: string[];
    /**
     * Fallback model to use when the primary model returns a rate-limit error.
     * Retried once automatically.
     */
    fallbackModel?: string;
}

interface ToolEvent {
    tool: string;
    args: string;
    /** Matches the producing tool_execution_start/update/end event. */
    toolCallId?: string;
    /**
     * "running" while between tool_execution_start and tool_execution_end; flipped
     * to "done" on end. We store every in-flight call in recentTools (keyed by
     * toolCallId) rather than a single current-tool slot, because pi-agent-core
     * dispatches a turn's tool calls in parallel via Promise.all — a single slot
     * would let the second start overwrite the first.
     */
    status: "running" | "done";
    /**
     * Live progress of subagents spawned by this tool call. Populated only for
     * `subagent` tool calls, from the `partialResult.details.results` payload of
     * `tool_execution_update` events (and refreshed once more from the end
     * event's final results). Recursive: each child's own progress may carry
     * further children via its `recentTools[i].children`.
     */
    children?: AgentResult[];
}

interface AgentProgress {
    agent: string;
    status: "pending" | "running" | "completed" | "failed";
    task: string;
    /**
     * Chronological log of tool calls — running and done interleaved. The
     * renderer prefixes running entries with `▸` and done ones with ` `.
     */
    recentTools: ToolEvent[];
    toolCount: number;
    tokens: number;
    durationMs: number;
    lastMessage: string;
    error?: string;
}

interface AgentResult {
    agent: string;
    task: string;
    output: string;
    exitCode: number;
    progress: AgentProgress;
    model?: string;
    contextWindow?: number;
    usage: { input: number; output: number; cacheRead: number; cacheWrite: number; cost: number; turns: number };
}

interface Details {
    results: AgentResult[];
}

// ── Config ─────────────────────────────────────────────────────────────────────

interface ExtensionConfig {
    maxConcurrency?: number;
}

const EXT_DIR = path.dirname(fileURLToPath(import.meta.url));
const DEFAULT_AGENTS_DIR = path.join(EXT_DIR, "agents");
const TOOLS_DIR = path.join(EXT_DIR, "tools");
const CONFIG_PATH = path.join(EXT_DIR, "config.json");
const DEFAULT_MAX_CONCURRENCY = 4;

function loadConfig(): ExtensionConfig {
    try {
        if (fs.existsSync(CONFIG_PATH)) {
            return JSON.parse(fs.readFileSync(CONFIG_PATH, "utf-8")) as ExtensionConfig;
        }
    } catch {}
    return {};
}

// Built-in tools that pi provides natively (no extension needed)
const BUILTIN_TOOLS = new Set(["read", "write", "edit", "bash", "grep", "find", "ls"]);

// Custom tools that require loading an extension into the subagent process
const EXT_BASE = agentPath("extensions");
const HERMES_MEMORY_EXTENSION = agentPath("npm", "node_modules", "pi-hermes-memory", "src", "index.ts");

const CUSTOM_TOOL_EXTENSIONS: Record<string, string> = {
    memory_search: HERMES_MEMORY_EXTENSION,
    session_search: HERMES_MEMORY_EXTENSION,
    web_search: path.join(EXT_BASE, "web-search", "index.ts"),
    web_fetch: path.join(EXT_BASE, "web-fetch", "index.ts"),
    safe_bash: path.join(TOOLS_DIR, "safe-bash.ts"),
    video_extract: path.join(EXT_BASE, "video-extract", "index.ts"),
    youtube_search: path.join(EXT_BASE, "youtube-search", "index.ts"),
    google_image_search: path.join(EXT_BASE, "google-image-search", "index.ts"),
    // `subagent` is the tool this very extension registers. Listing it here lets
    // a parent agent grant it to a child agent — the child pi process loads this
    // same index.ts via `--extension`, sees its own subagent tool, and (if
    // PI_SUBAGENT_ALLOWED is set) only registers the allowlisted agents.
    subagent: path.join(EXT_DIR, "index.ts"),
};

// ── Agent Discovery & Registration ────────────────────────────────────────────

let agents: AgentConfig[] = [];
let runtimeAllowedAgents: string[] | undefined;

// Read once at module load. If we're a child subagent process whose parent
// pinned an allowlist, we silently ignore any agent (built-in OR registered
// later by a third-party extension) that isn't in the list.
const SUBAGENT_ALLOWLIST: string[] | undefined = (() => {
    const raw = process.env.PI_SUBAGENT_ALLOWED;
    if (!raw) return undefined;
    const list = raw.split(",").map((s) => s.trim()).filter(Boolean);
    return list.length > 0 ? list : undefined;
})();

export function registerAgent(config: AgentConfig): void {
    if (SUBAGENT_ALLOWLIST && !SUBAGENT_ALLOWLIST.includes(config.name)) return;
    if (agents.find((a) => a.name === config.name)) {
        throw new Error(`Agent already registered: ${config.name}`);
    }
    agents.push(config);
}

export function unregisterAgent(name: string): void {
    agents = agents.filter((a) => a.name !== name);
}

export function setAllowedAgents(names: string[] | undefined): void {
    runtimeAllowedAgents = names?.length ? names.map((name) => name.toLowerCase()) : undefined;
}

function allowedBackgroundAgents(): AgentConfig[] {
    return agents.filter((agent) => {
        if (agent.role === "foreground") return false;
        if (!runtimeAllowedAgents) return true;
        return runtimeAllowedAgents.includes(agent.name.toLowerCase());
    });
}

// Expose registration functions globally so other extensions loaded via jiti
// (which creates separate module instances) can access the shared agents array.
(globalThis as any).__pi_subagents = { registerAgent, unregisterAgent, setAllowedAgents };

function loadAgents(cwd = process.cwd()): AgentConfig[] {
    return loadAgentDefinitions({ defaultAgentsDir: DEFAULT_AGENTS_DIR, cwd }).map((def) => ({
        name: def.name,
        description: def.description,
        tools: def.tools,
        // Agents without an explicit model inherit the foreground model at execution time.
        model: def.model || "",
        thinking: def.thinking || "medium",
        systemPrompt: def.systemPrompt,
        filePath: def.filePath,
        role: def.role,
        subagentAgents: def.subagentAgents,
        fallbackModel: def.fallbackModel,
    } satisfies AgentConfig));
}

// ── Pi Binary Resolution ───────────────────────────────────────────────────────

function resolvePiBinary(): { command: string; baseArgs: string[] } {
    // Resolve the pi entry point from process.argv[1]
    const entry = process.argv[1];
    if (entry) {
        try {
            const realEntry = fs.realpathSync(entry);
            if (/\.(?:mjs|cjs|js)$/i.test(realEntry)) {
                return { command: process.execPath, baseArgs: [realEntry] };
            }
        } catch {}
    }
    return { command: process.platform === "win32" ? "pi.cmd" : "pi", baseArgs: [] };
}

// ── Formatting Utilities ───────────────────────────────────────────────────────

function formatTokens(n: number): string {
    return n < 1000 ? String(n) : n < 10000 ? `${(n / 1000).toFixed(1)}k` : `${Math.round(n / 1000)}k`;
}

function formatDuration(ms: number): string {
    if (ms < 1000) return `${ms}ms`;
    if (ms < 60000) return `${(ms / 1000).toFixed(1)}s`;
    return `${Math.floor(ms / 60000)}m${Math.floor((ms % 60000) / 1000)}s`;
}

function formatContextUsage(tokens: number, contextWindow: number | undefined): string {
    if (!contextWindow) return `${formatTokens(tokens)} ctx`;
    const pct = (tokens / contextWindow) * 100;
    const maxStr = contextWindow >= 1_000_000
        ? `${(contextWindow / 1_000_000).toFixed(1)}M`
        : `${Math.round(contextWindow / 1000)}k`;
    return `${pct.toFixed(1)}%/${maxStr}`;
}

function truncLine(text: string, maxWidth: number): string {
    const flat = text.replace(/\r?\n/g, "↵ ");
    return truncateToWidth(flat, Math.max(1, maxWidth));
}

function wrapAnsiLines(text: string, maxWidth: number): string[] {
    return wrapTextWithAnsi(text, Math.max(1, maxWidth));
}

function boundedAppendTail(current: string, chunk: string, maxBytes = DEFAULT_MAX_BYTES): string {
    const combined = current + chunk;
    if (Buffer.byteLength(combined, "utf8") <= maxBytes) return combined;
    const prefix = `[truncated to last ${maxBytes} bytes]\n`;
    const trunc = truncateTail(combined, {
        maxLines: DEFAULT_MAX_LINES,
        maxBytes: maxBytes - Buffer.byteLength(prefix, "utf8"),
    });
    return trunc.truncated ? `${prefix}${trunc.content}` : trunc.content;
}

function trimPendingLine(current: string, maxBytes = DEFAULT_MAX_BYTES): string {
    if (Buffer.byteLength(current, "utf8") <= maxBytes) return current;
    const trunc = truncateTail(current, { maxLines: DEFAULT_MAX_LINES, maxBytes });
    return trunc.content;
}

// ── Subagent Execution ─────────────────────────────────────────────────────────

async function buildPiArgs(
    agent: AgentConfig,
    task: string,
    cwd: string,
): Promise<{ args: string[]; tempDir: string; childEnv: NodeJS.ProcessEnv | undefined; missingTools: string[] }> {
    const piBin = resolvePiBinary();
    const tempDir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-sub-"));

    // Write system prompt to temp file
    const promptPath = path.join(tempDir, `${agent.name}.md`);
    await withFileMutationQueue(promptPath, async () => {
        await fs.promises.writeFile(promptPath, agent.systemPrompt, { encoding: "utf-8", mode: 0o600 });
    });

    const args = [...piBin.baseArgs, "--mode", "json", "-p", "--no-session", "--no-skills"];

    const allowlist: string[] = [];
    const extensionPaths = new Set<string>();
    const missingTools: string[] = [];
    for (const tool of agent.tools) {
        if (BUILTIN_TOOLS.has(tool)) {
            allowlist.push(tool);
        } else if (CUSTOM_TOOL_EXTENSIONS[tool]) {
            const extensionPath = CUSTOM_TOOL_EXTENSIONS[tool];
            if (fs.existsSync(extensionPath)) {
                allowlist.push(tool);
                extensionPaths.add(extensionPath);
            } else {
                missingTools.push(`${tool} (extension not found: ${extensionPath})`);
            }
        } else {
            missingTools.push(`${tool} (unknown tool)`);
        }
    }

    args.push("--no-extensions");
    if (allowlist.length > 0) {
        args.push("--tools", allowlist.join(","));
    } else {
        args.push("--no-tools");
    }
    for (const extPath of extensionPaths) {
        args.push("--extension", extPath);
    }

    args.push("--model", agent.model);
    args.push("--thinking", agent.thinking);
    args.push("--append-system-prompt", promptPath);

    // Keep long or shell-sensitive Windows fallback tasks out of cmd.exe's
    // command line. Pi expands @file arguments before model execution.
    const TASK_LIMIT = 8000;
    const shellFallback = process.platform === "win32" && /\.(?:cmd|bat)$/i.test(piBin.command);
    if (task.length > TASK_LIMIT || shellFallback) {
        const taskPath = path.join(tempDir, "task.md");
        await withFileMutationQueue(taskPath, async () => {
            await fs.promises.writeFile(taskPath, `Task: ${task}`, { encoding: "utf-8", mode: 0o600 });
        });
        args.push(`@${taskPath}`);
    } else {
        args.push(`Task: ${task}`);
    }

    // If this agent is allowed to spawn subagents AND we want to restrict which
    // ones, pass the allowlist down via env.
    let childEnv: NodeJS.ProcessEnv | undefined;
    if (agent.tools.includes("subagent") && agent.subagentAgents && agent.subagentAgents.length > 0) {
        childEnv = { ...process.env, PI_SUBAGENT_ALLOWED: agent.subagentAgents.join(",") };
    }

    return { args: [piBin.command, ...args], tempDir, childEnv, missingTools };
}

function extractTextFromContent(content: unknown): string {
    if (!content) return "";
    if (typeof content === "string") return content;
    if (Array.isArray(content)) {
        return content
            .filter((c: any) => c.type === "text")
            .map((c: any) => c.text)
            .join("\n");
    }
    return "";
}

/** Collapse any whitespace run (incl. newlines) into a single space. */
function flatten(s: string): string {
    return s.replace(/\s+/g, " ").trim();
}

const MAX_ARG_PREVIEW = 4000;

function extractToolArgsPreview(args: Record<string, unknown>): string {
    const cap = (s: string) => (s.length > MAX_ARG_PREVIEW ? s.slice(0, MAX_ARG_PREVIEW) + "…" : s);
    if (args.command) return cap(flatten(String(args.command)));
    if (args.path) return cap(flatten(String(args.path)));
    if (args.query) return `"${cap(flatten(String(args.query)))}"`;
    if (args.url) return cap(flatten(String(args.url)));
    if (args.pattern) return cap(flatten(String(args.pattern)));
    if (args.agent) return flatten(String(args.agent));
    if (Array.isArray(args.tasks)) {
        const names = (args.tasks as Array<{ agent?: string }>)
            .map((t) => t?.agent || "?")
            .join(", ");
        return `parallel(${names})`;
    }
    return cap(flatten(JSON.stringify(args)));
}

async function runSubagent(
    agent: AgentConfig,
    task: string,
    cwd: string,
    signal: AbortSignal | undefined,
    onUpdate?: (progress: AgentProgress, usage: AgentResult["usage"]) => void,
): Promise<AgentResult> {
    const { args, tempDir, childEnv, missingTools } = await buildPiArgs(agent, task, cwd);
    if (missingTools.length > 0) {
        try { fs.rmSync(tempDir, { recursive: true, force: true }); } catch {}
        throw new Error(
            `Agent "${agent.name}" requested unavailable tools: ${missingTools.join("; ")}. ` +
            `Check that required extensions are installed and paths are correct.`,
        );
    }
    const command = args[0];
    const spawnArgs = args.slice(1);

    const result: AgentResult = {
        agent: agent.name,
        task,
        output: "",
        exitCode: 0,
        model: agent.model,
        usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0 },
        progress: {
            agent: agent.name,
            status: "running",
            task,
            recentTools: [],
            toolCount: 0,
            tokens: 0,
            durationMs: 0,
            lastMessage: "",
        },
    };

    const startTime = Date.now();
    const progress = result.progress;
    const fireUpdate = throttle(() => {
        progress.durationMs = Date.now() - startTime;
        onUpdate?.(progress, result.usage);
    }, 150);

    // .cmd/.bat files on Windows need shell: true for reliable argument passing;
    // CreateProcess alone can mishandle quoting for complex command lines.
    const isWinCmd = process.platform === "win32" && /\.(?:cmd|bat)$/i.test(command);
    const spawnOpts: any = {
        cwd,
        stdio: ["ignore", "pipe", "pipe"] as const,
        ...(childEnv ? { env: childEnv } : {}),
        ...(process.platform !== "win32" ? { detached: true } : {}),
        ...(isWinCmd ? { shell: true } : {}),
    };

    const exitCode = await new Promise<number>((resolve) => {
        const proc = spawn(command, spawnArgs, spawnOpts);

        let buf = "";
        let stderrBuf = "";
        let forcedKillTimer: ReturnType<typeof setTimeout> | undefined;
        let aborted = false;

        const killProcessTree = (force = false) => {
            if (!proc.pid) return;
            if (process.platform === "win32") {
                const args = ["/pid", String(proc.pid), "/t"];
                if (force) args.push("/f");
                try {
                    const killer = spawn("taskkill", args, { stdio: "ignore", windowsHide: true });
                    killer.on("error", () => {});
                } catch {
                    try { proc.kill(force ? "SIGKILL" : "SIGTERM"); } catch {}
                }
                return;
            }
            try { process.kill(-proc.pid, force ? "SIGKILL" : "SIGTERM"); }
            catch { try { proc.kill(force ? "SIGKILL" : "SIGTERM"); } catch {} }
        };

        const onAbort = () => {
            aborted = true;
            killProcessTree(false);
            forcedKillTimer = setTimeout(() => killProcessTree(true), 3000);
        };

        const cleanup = () => {
            if (forcedKillTimer) { clearTimeout(forcedKillTimer); forcedKillTimer = undefined; }
            signal?.removeEventListener("abort", onAbort);
        };

        const processLine = (line: string) => {
            if (!line.trim()) return;
            try {
                const evt = JSON.parse(line) as any;
                progress.durationMs = Date.now() - startTime;

                if (evt.type === "tool_execution_start") {
                    progress.toolCount++;
                    progress.recentTools.push({
                        tool: evt.toolName,
                        args: extractToolArgsPreview((evt.args || {}) as Record<string, unknown>),
                        toolCallId: evt.toolCallId,
                        status: "running",
                    });
                    fireUpdate();
                }

                if (evt.type === "tool_execution_update") {
                    const partial = evt.partialResult as { details?: { results?: unknown } } | undefined;
                    const nested = partial?.details?.results;
                    if (evt.toolName === "subagent" && Array.isArray(nested) && evt.toolCallId) {
                        const hit = progress.recentTools.find((t) => t.toolCallId === evt.toolCallId);
                        if (hit) { hit.children = nested as AgentResult[]; fireUpdate(); }
                    }
                }

                if (evt.type === "tool_execution_end") {
                    const hit = evt.toolCallId
                        ? progress.recentTools.find((t) => t.toolCallId === evt.toolCallId)
                        : undefined;
                    if (hit) {
                        hit.status = "done";
                        const finalResult = evt.result as { details?: { results?: unknown } } | undefined;
                        const finalChildren = finalResult?.details?.results;
                        if (evt.toolName === "subagent" && Array.isArray(finalChildren)) {
                            hit.children = finalChildren as AgentResult[];
                        }
                    }
                    fireUpdate();
                }

                if (evt.type === "tool_result_end") {
                    fireUpdate();
                }

                if (evt.type === "message_end" && evt.message) {
                    if (evt.message.role === "assistant") {
                        result.usage.turns++;
                        const u = evt.message.usage;
                        if (u) {
                            result.usage.input += u.input || 0;
                            result.usage.output += u.output || 0;
                            result.usage.cacheRead += u.cacheRead || 0;
                            result.usage.cacheWrite += u.cacheWrite || 0;
                            result.usage.cost += u.cost?.total || 0;
                            progress.tokens = (u as { totalTokens?: number }).totalTokens ||
                                (u.input || 0) + (u.output || 0) + (u.cacheRead || 0) + (u.cacheWrite || 0);
                        }
                        if (evt.message.model) result.model = evt.message.model;
                        if (evt.message.errorMessage) progress.error = evt.message.errorMessage;
                        const text = extractTextFromContent(evt.message.content);
                        if (text) {
                            result.output = text;
                            const proseLines: string[] = [];
                            let inCodeBlock = false;
                            for (const line of text.split("\n")) {
                                if (line.trimStart().startsWith("```")) { inCodeBlock = !inCodeBlock; continue; }
                                if (!inCodeBlock && line.trim()) proseLines.push(line.trim());
                            }
                            if (proseLines.length > 0) progress.lastMessage = proseLines.slice(0, 3).join(" ");
                        }
                    }
                    fireUpdate();
                }
            } catch {
                // Non-JSON lines are expected
            }
        };

        proc.stdout!.on("data", (d: Buffer) => {
            buf = boundedAppendTail(buf, d.toString());
            const lines = buf.split("\n");
            buf = trimPendingLine(lines.pop() || "");
            lines.forEach(processLine);
        });

        proc.stderr!.on("data", (d: Buffer) => { stderrBuf = boundedAppendTail(stderrBuf, d.toString()); });

        proc.on("close", (code) => {
            cleanup();
            if (buf.trim()) processLine(buf);
            if (aborted && !progress.error) progress.error = "Operation aborted";
            if (code !== 0 && stderrBuf.trim() && !progress.error) {
                progress.error = stderrBuf.trim();
            }
            resolve(code ?? 1);
        });

        proc.on("error", () => { cleanup(); resolve(1); });

        if (signal) {
            if (signal.aborted) onAbort();
            else signal.addEventListener("abort", onAbort, { once: true });
        }
    });

    // Cleanup temp dir
    try { fs.rmSync(tempDir, { recursive: true, force: true }); } catch {}

    result.exitCode = exitCode;
    progress.status = exitCode === 0 && !progress.error ? "completed" : "failed";
    progress.durationMs = Date.now() - startTime;
    if (progress.error) result.output = result.output || `Error: ${progress.error}`;

    if (result.output.length > DEFAULT_MAX_BYTES) {
        const trunc = truncateHead(result.output, { maxLines: DEFAULT_MAX_LINES, maxBytes: DEFAULT_MAX_BYTES });
        result.output = trunc.content;
        if (trunc.truncated) result.output += "\n\n[Output truncated]";
    }

    return result;
}

// ── Rate-limit detection ──────────────────────────────────────────────────────

function isRateLimitError(error: string | undefined): boolean {
    if (!error) return false;
    return /429|rate.?limit|too many requests|quota exceeded/i.test(error);
}

function formatFailedAgentError(result: AgentResult): string {
    const output = result.output || result.progress.error || "No output";
    return `[${result.agent}] failed with exit code ${result.exitCode}\n\n${output}`;
}

// ── Throttle ───────────────────────────────────────────────────────────────────

function throttle<T extends (...args: any[]) => void>(fn: T, ms: number): T {
    let lastCall = 0;
    let timer: ReturnType<typeof setTimeout> | undefined;
    return ((...args: any[]) => {
        const now = Date.now();
        const remaining = ms - (now - lastCall);
        if (remaining <= 0) {
            lastCall = now;
            if (timer) { clearTimeout(timer); timer = undefined; }
            fn(...args);
        } else if (!timer) {
            timer = setTimeout(() => {
                lastCall = Date.now();
                timer = undefined;
                fn(...args);
            }, remaining);
        }
    }) as T;
}

// ── Parallel Execution with Concurrency Limit ──────────────────────────────────

class Semaphore {
    private inFlight = 0;
    private readonly waiters: Array<() => void> = [];
    constructor(private readonly max: number) {}
    async run<T>(fn: () => Promise<T>): Promise<T> {
        if (this.inFlight >= this.max) {
            await new Promise<void>((r) => this.waiters.push(r));
        }
        this.inFlight++;
        try {
            return await fn();
        } finally {
            this.inFlight--;
            const next = this.waiters.shift();
            if (next) next();
        }
    }
}

// ── Rendering ──────────────────────────────────────────────────────────────────

type Theme = ExtensionContext["ui"]["theme"];


function renderAgentProgress(
    r: AgentResult,
    theme: Theme,
    expanded: boolean,
    w: number,
    depth: number = 0,
): Container {
    const c = new Container();
    const prog = r.progress;
    const isRunning = prog.status === "running";
    const isPending = prog.status === "pending";
    const nested = depth > 0;

    const indent = nested
        ? "  ".repeat(Math.min(depth, Math.max(0, Math.floor((Math.max(1, w) - 1) / 2))))
        : "";
    const innerW = Math.max(1, w - visibleWidth(indent));

    const addLine = (content: string) => {
        if (expanded) {
            for (const line of wrapAnsiLines(content, innerW)) {
                c.addChild(new Text(indent + line, 0, 0));
            }
        } else {
            c.addChild(new Text(indent + truncLine(content, innerW), 0, 0));
        }
    };

    const icon = isRunning
        ? theme.fg("warning", "⟳")
        : isPending
        ? theme.fg("dim", "○")
        : r.exitCode === 0
        ? theme.fg("success", "✓")
        : theme.fg("error", "✗");

    const stats = `${prog.toolCount} tools · ${formatDuration(prog.durationMs)}`;
    const modelStr = r.model ? theme.fg("dim", ` (${r.model})`) : "";
    addLine(`${icon} ${theme.fg("toolTitle", theme.bold(r.agent))}${modelStr} — ${theme.fg("dim", stats)}`);

    const renderToolRow = (
        toolName: string,
        args: string,
        children: AgentResult[] | undefined,
        isCurrent: boolean,
    ) => {
        const body = args ? `${toolName}: ${args}` : toolName;
        if (isCurrent) {
            addLine(theme.fg("warning", `▸ ${body}`));
        } else {
            addLine(theme.fg("muted", `  ${body}`));
        }
        if (children && children.length > 0) {
            for (const child of children) {
                c.addChild(renderAgentProgress(child, theme, expanded, w, depth + 1));
            }
        }
    };

    for (const t of prog.recentTools) {
        renderToolRow(t.tool, t.args, t.children, t.status === "running");
    }

    if (prog.lastMessage) {
        if (!nested) c.addChild(new Spacer(1));
        addLine(theme.fg("text", prog.lastMessage));
    }

    if (!nested && !isRunning && r.output && expanded) {
        c.addChild(new Spacer(1));
        const mdTheme = getMarkdownTheme();
        c.addChild(new Markdown(r.output, 0, 0, mdTheme));
    }

    if (!nested) c.addChild(new Spacer(1));

    const usageParts: string[] = [];
    if (r.usage.input) usageParts.push(theme.fg("dim", `↑${formatTokens(r.usage.input)}`));
    if (r.usage.output) usageParts.push(theme.fg("dim", `↓${formatTokens(r.usage.output)}`));
    if (r.usage.cacheRead) usageParts.push(theme.fg("dim", `R${formatTokens(r.usage.cacheRead)}`));
    if (r.usage.cacheWrite) usageParts.push(theme.fg("dim", `W${formatTokens(r.usage.cacheWrite)}`));
    if (r.usage.cost) usageParts.push(theme.fg("dim", `$${r.usage.cost.toFixed(3)}`));
    if (prog.tokens > 0) {
        const ctxStr = formatContextUsage(prog.tokens, r.contextWindow);
        const pct = r.contextWindow ? (prog.tokens / r.contextWindow) * 100 : 0;
        const coloredCtx = pct > 90
            ? theme.fg("error", ctxStr)
            : pct > 70
            ? theme.fg("warning", ctxStr)
            : theme.fg("dim", ctxStr);
        usageParts.push(coloredCtx);
    }
    if (usageParts.length) addLine(usageParts.join(" "));
    if (prog.error) addLine(theme.fg("error", `Error: ${prog.error}`));

    return c;
}

class WrappedText implements Component {
    constructor(private readonly text: string, private readonly indentWidth = 0) {}
    render(width: number): string[] {
        const inner = Math.max(1, width - this.indentWidth);
        return wrapAnsiLines(this.text, inner);
    }
    invalidate(): void {}
}

class WidthAwareLine implements Component {
    constructor(private readonly build: (width: number) => string) {}
    render(width: number): string[] { return [this.build(width)]; }
    invalidate(): void {}
}

class AgentProgressComponent implements Component {
    constructor(
        private readonly result: AgentResult,
        private readonly theme: Theme,
        private readonly expanded: boolean,
    ) {}
    render(width: number): string[] {
        return renderAgentProgress(this.result, this.theme, this.expanded, Math.max(1, width)).render(width);
    }
    invalidate(): void {}
}

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

export default function (pi: ExtensionAPI) {
    const config = loadConfig();
    const semaphore = new Semaphore(config.maxConcurrency ?? DEFAULT_MAX_CONCURRENCY);
    agents = loadAgents();

    // If spawned as a child by a parent subagent process, PI_SUBAGENT_ALLOWED
    // pins which agents we're allowed to expose. Filter the registry now.
    if (SUBAGENT_ALLOWLIST) {
        agents = agents.filter((a) => SUBAGENT_ALLOWLIST!.includes(a.name));
    }

    pi.registerTool({
        name: "subagent",
        label: "Subagent",
        description:
            "Run a subagent to complete a task. Subagents have NO context from the current conversation — include all necessary context in the task description.",
        promptSnippet: "Run subagents for delegated tasks",
        promptGuidelines: [
            "Parallel tool calls are your primary parallelism mechanism — put multiple independent read/fetch/search calls in one function_calls block. Don't use subagents to parallelize simple I/O.",
            "Use subagent to delegate *reasoning and decisions*: exploration/research (explorer), review (critic), advice (advisor), or isolated code changes (coder)",
            "For multiple independent subagent tasks, emit multiple `subagent` tool calls in the same turn — they run in parallel automatically.",
            "Subagents have NO context from the current conversation — include ALL necessary context in the task description",
        ],
        parameters: Type.Object({
            agent: Type.String({ description: "Name of the agent to invoke" }),
            task: Type.String({ description: "Task description" }),
            cwd: Type.Optional(Type.String({ description: "Working directory for the agent process" })),
        }),
        async execute(toolCallId, params, signal, onUpdate, ctx) {
            const cwd = ctx.cwd;

            if (!params.agent || !params.task) {
                throw new Error(
                    "`subagent` requires both `agent` and `task`. To fan out work, emit multiple `subagent` tool calls in the same turn — they run in parallel.",
                );
            }

            const availableAgents = allowedBackgroundAgents();
            const configuredAgent = availableAgents.find((a) => a.name.toLowerCase() === params.agent.toLowerCase());
            if (!configuredAgent) {
                const available = availableAgents.map((a) => a.name).join(", ") || "none";
                throw new Error(`Unknown or unavailable background agent: ${params.agent}. Available agents: ${available}`);
            }

            const inheritedModel = ctx.model?.provider && ctx.model.id
                ? `${ctx.model.provider}/${ctx.model.id}`
                : "openai-codex/gpt-5.6-terra";
            const agent = configuredAgent.model
                ? configuredAgent
                : { ...configuredAgent, model: inheritedModel };
            const [provider, modelId] = agent.model.split("/");
            const contextWindow =
                provider && modelId
                    ? ctx.modelRegistry.find(provider, modelId)?.contextWindow
                    : undefined;

            const liveResult: AgentResult = {
                agent: params.agent,
                task: params.task,
                output: "",
                exitCode: -1,
                model: agent.model,
                contextWindow,
                usage: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0 },
                progress: {
                    agent: params.agent,
                    status: "running" as const,
                    task: params.task,
                    recentTools: [],
                    toolCount: 0,
                    tokens: 0,
                    durationMs: 0,
                    lastMessage: "",
                },
            };

            let result = await semaphore.run(() =>
                runSubagent(agent, params.task!, params.cwd ?? cwd, signal, (progress, usage) => {
                    liveResult.progress = progress;
                    liveResult.usage = { ...usage };
                    onUpdate?.({
                        content: [{ type: "text", text: "(running...)" }],
                        details: { results: [liveResult] },
                    });
                }),
            );

            // Fallback model retry on rate limit
            if (result.exitCode !== 0 && agent.fallbackModel && isRateLimitError(result.progress.error)) {
                const fallbackAgent = { ...agent, model: agent.fallbackModel };
                liveResult.model = agent.fallbackModel;
                liveResult.progress = {
                    ...liveResult.progress,
                    status: "running" as const,
                    error: undefined,
                    recentTools: [],
                    toolCount: 0,
                    tokens: 0,
                    durationMs: 0,
                    lastMessage: `Rate limited on ${agent.model}, retrying with ${agent.fallbackModel}…`,
                };
                onUpdate?.({
                    content: [{ type: "text", text: `(rate limited, retrying with ${agent.fallbackModel}…)` }],
                    details: { results: [liveResult] },
                });
                result = await semaphore.run(() =>
                    runSubagent(fallbackAgent, params.task!, params.cwd ?? cwd, signal, (progress, usage) => {
                        liveResult.progress = progress;
                        liveResult.usage = { ...usage };
                        onUpdate?.({
                            content: [{ type: "text", text: "(running...)" }],
                            details: { results: [liveResult] },
                        });
                    }),
                );
            }

            result.contextWindow = contextWindow;
            const isError = result.exitCode !== 0 || !!result.progress.error;
            if (isError) throw new Error(formatFailedAgentError(result));

            return {
                content: [{ type: "text", text: result.output || "(no output)" }],
                details: { results: [result] },
            };
        },

        // ── Render: tool call header ──
        renderCall(args, theme, context) {
            if (!context.expanded) {
                return new WidthAwareLine((width) => {
                    if (!args.agent) return truncLine(theme.fg("toolTitle", theme.bold("subagent")), width);
                    const prefix = `${theme.fg("toolTitle", theme.bold("subagent"))} ${theme.fg("accent", args.agent)} `;
                    if (visibleWidth(prefix) >= width) return truncLine(prefix, width);
                    const taskPreview = args.task ? args.task.replace(/\r?\n/g, " ") : "";
                    return prefix + theme.fg("dim", truncLine(taskPreview, width - visibleWidth(prefix)));
                });
            }

            const c =
                context.lastComponent instanceof Container
                    ? (context.lastComponent.clear(), context.lastComponent)
                    : new Container();
            const agentLabel = args.agent ? ` ${theme.fg("accent", args.agent)}` : "";
            const cwdLabel = args.cwd ? theme.fg("dim", ` (cwd: ${args.cwd})`) : "";
            c.addChild(new Text(`${theme.fg("toolTitle", theme.bold("subagent"))}${agentLabel}${cwdLabel}`, 0, 0));
            if (args.task) {
                c.addChild(new Spacer(1));
                c.addChild(new WrappedText(theme.fg("text", args.task), 0));
            }
            return c;
        },

        // ── Render: result ──
        renderResult(result, options, theme, context) {
            const details = result.details as Details | undefined;
            if (!details?.results?.length) {
                const t = result.content[0];
                const text = t?.type === "text" ? t.text : "(no output)";
                return new Text(text.slice(0, 200), 0, 0);
            }
            return new AgentProgressComponent(details.results[0], theme, options.expanded);
        },
    });
}