/**
 * Subagent Tool - Delegate tasks to specialized agents.
 *
 * Creates in-process SDK sessions for each subagent invocation,
 * giving it an isolated context window.
 *
 * Supports two modes:
 *   - Single:   { agent: "name", task: "..." }
 *   - Parallel: { tasks: [{ agent: "name", task: "..." }, ...] }
 *
 * Sessions are persisted to disk for observability.
 */

import * as os from "node:os";
import type { Message } from "@earendil-works/pi-ai";
import { type ExtensionAPI, getMarkdownTheme } from "@earendil-works/pi-coding-agent";
import { Container, Markdown, Spacer, Text } from "@earendil-works/pi-tui";
import { Type } from "typebox";
import { type AgentConfig, discoverAgents, formatAgentList } from "./agent.js";
import {
    type SingleResult,
    runAgent,
    runParallel,
} from "./runner.js";

const MAX_PARALLEL_TASKS = 8;
const PER_TASK_OUTPUT_CAP = 50 * 1024;
const COLLAPSED_ITEM_COUNT = 10;

function formatTokens(count: number): string {
    if (count < 1000) return count.toString();
    if (count < 10000) return `${(count / 1000).toFixed(1)}k`;
    if (count < 1000000) return `${Math.round(count / 1000)}k`;
    return `${(count / 1000000).toFixed(1)}M`;
}

function formatUsageStats(
    usage: {
        input: number;
        output: number;
        cacheRead: number;
        cacheWrite: number;
        cost: number;
        contextTokens?: number;
        turns?: number;
    },
    model?: string,
): string {
    const parts: string[] = [];
    if (usage.turns) parts.push(`${usage.turns} turn${usage.turns > 1 ? "s" : ""}`);
    if (usage.input) parts.push(`↑${formatTokens(usage.input)}`);
    if (usage.output) parts.push(`↓${formatTokens(usage.output)}`);
    if (usage.cacheRead) parts.push(`R${formatTokens(usage.cacheRead)}`);
    if (usage.cacheWrite) parts.push(`W${formatTokens(usage.cacheWrite)}`);
    if (usage.cost) parts.push(`$${usage.cost.toFixed(4)}`);
    if (usage.contextTokens && usage.contextTokens > 0) {
        parts.push(`ctx:${formatTokens(usage.contextTokens)}`);
    }
    if (model) parts.push(model);
    return parts.join(" ");
}

function formatToolCall(
    toolName: string,
    args: Record<string, unknown>,
    themeFg: (color: any, text: string) => string,
): string {
    const shortenPath = (p: string) => {
        const home = os.homedir();
        return p.startsWith(home) ? `~${p.slice(home.length)}` : p;
    };

    switch (toolName) {
        case "bash": {
            const command = (args.command as string) || "...";
            const preview = command.length > 60 ? `${command.slice(0, 60)}...` : command;
            return themeFg("muted", "$ ") + themeFg("toolOutput", preview);
        }
        case "read": {
            const rawPath = (args.file_path || args.path || "...") as string;
            const filePath = shortenPath(rawPath);
            const offset = args.offset as number | undefined;
            const limit = args.limit as number | undefined;
            let text = themeFg("accent", filePath);
            if (offset !== undefined || limit !== undefined) {
                const startLine = offset ?? 1;
                const endLine = limit !== undefined ? startLine + limit - 1 : "";
                text += themeFg("warning", `:${startLine}${endLine ? `-${endLine}` : ""}`);
            }
            return themeFg("muted", "read ") + text;
        }
        case "write": {
            const rawPath = (args.file_path || args.path || "...") as string;
            const filePath = shortenPath(rawPath);
            const content = (args.content || "") as string;
            const lines = content.split("\n").length;
            let text = themeFg("muted", "write ") + themeFg("accent", filePath);
            if (lines > 1) text += themeFg("dim", ` (${lines} lines)`);
            return text;
        }
        case "edit": {
            const rawPath = (args.file_path || args.path || "...") as string;
            return themeFg("muted", "edit ") + themeFg("accent", shortenPath(rawPath));
        }
        case "ls": {
            const rawPath = (args.path || ".") as string;
            return themeFg("muted", "ls ") + themeFg("accent", shortenPath(rawPath));
        }
        case "find": {
            const pattern = (args.pattern || "*") as string;
            const rawPath = (args.path || ".") as string;
            return (
                themeFg("muted", "find ") +
                themeFg("accent", pattern) +
                themeFg("dim", ` in ${shortenPath(rawPath)}`)
            );
        }
        case "grep": {
            const pattern = (args.pattern || "") as string;
            const rawPath = (args.path || ".") as string;
            return (
                themeFg("muted", "grep ") +
                themeFg("accent", `/${pattern}/`) +
                themeFg("dim", ` in ${shortenPath(rawPath)}`)
            );
        }
        default: {
            const argsStr = JSON.stringify(args);
            const preview = argsStr.length > 50 ? `${argsStr.slice(0, 50)}...` : argsStr;
            return themeFg("accent", toolName) + themeFg("dim", ` ${preview}`);
        }
    }
}

interface SubagentDetails {
    mode: "single" | "parallel";
    results: SingleResult[];
}

function getFinalOutput(messages: Message[]): string {
    for (let i = messages.length - 1; i >= 0; i--) {
        const msg = messages[i];
        if (msg.role === "assistant") {
            for (const part of msg.content) {
                if (part.type === "text") return part.text;
            }
        }
    }
    return "";
}

function isFailedResult(result: SingleResult): boolean {
    return result.exitCode !== 0 || result.stopReason === "error" || result.stopReason === "aborted";
}

function getResultOutput(result: SingleResult): string {
    if (isFailedResult(result)) {
        return result.errorMessage || result.stderr || getFinalOutput(result.messages) || "(no output)";
    }
    return getFinalOutput(result.messages) || "(no output)";
}

function truncateParallelOutput(output: string): string {
    const byteLength = Buffer.byteLength(output, "utf8");
    if (byteLength <= PER_TASK_OUTPUT_CAP) return output;
    let truncated = output.slice(0, PER_TASK_OUTPUT_CAP);
    while (Buffer.byteLength(truncated, "utf8") > PER_TASK_OUTPUT_CAP) {
        truncated = truncated.slice(0, -1);
    }
    return `${truncated}\n\n[Output truncated: ${byteLength - Buffer.byteLength(truncated, "utf8")} bytes omitted. Full output preserved in tool details.]`;
}

type DisplayItem =
    | { type: "text"; text: string }
    | { type: "toolCall"; name: string; args: Record<string, any> };

function getDisplayItems(messages: Message[]): DisplayItem[] {
    const items: DisplayItem[] = [];
    for (const msg of messages) {
        if (msg.role === "assistant") {
            for (const part of msg.content) {
                if (part.type === "text") items.push({ type: "text", text: part.text });
                else if (part.type === "toolCall")
                    items.push({ type: "toolCall", name: part.name, args: part.arguments });
            }
        }
    }
    return items;
}

const TaskItem = Type.Object({
    agent: Type.String({ description: "Name of the agent to invoke" }),
    task: Type.String({ description: "Task to delegate to the agent" }),
    cwd: Type.Optional(Type.String({ description: "Working directory for the agent process" })),
});

const SubagentParams = Type.Object({
    agent: Type.Optional(Type.String({ description: "Name of the agent to invoke (for single mode)" })),
    task: Type.Optional(Type.String({ description: "Task to delegate (for single mode)" })),
    tasks: Type.Optional(
        Type.Array(TaskItem, { description: "Array of {agent, task} for parallel execution" }),
    ),
    cwd: Type.Optional(
        Type.String({ description: "Working directory for the agent process (single mode)" }),
    ),
});

const MAX_AGENTS_IN_DESCRIPTION = 20;

function buildDescription(agents: AgentConfig[]): string {
    const base =
        "Delegate tasks to specialized subagents with isolated context. " +
        "Modes: single (agent + task), parallel (tasks array).";

    if (agents.length === 0) {
        return `${base} No agents discovered.`;
    }

    const { text, remaining } = formatAgentList(agents, MAX_AGENTS_IN_DESCRIPTION);
    const suffix = remaining > 0 ? ` ... and ${remaining} more` : "";
    return `${base} Available agents: ${text}${suffix}`;
}

export default function (pi: ExtensionAPI) {
    let cachedAgents: AgentConfig[] = [];

    pi.on("session_start", async (_event, ctx) => {
        cachedAgents = discoverAgents(ctx);

        pi.registerCommand("agents", {
            description: "List available agents",
            handler: async (_args, ctx) => {
                if (cachedAgents.length === 0) {
                    ctx.ui.notify("No agents discovered.", "info");
                    return;
                }
                const lines = cachedAgents.map(
                    (a) => `${a.name} (${a.source}): ${a.description}`,
                );
                ctx.ui.notify(lines.join("\n"), "info");
            },
        });

        pi.registerTool({
            name: "delegate",
            label: "Delegate",
            description: buildDescription(cachedAgents),
            promptSnippet: "Delegate tasks to subagents",
            parameters: SubagentParams,

            async execute(_toolCallId, params, signal, onUpdate, ctx) {
                const agents = cachedAgents;

                const hasTasks = (params.tasks?.length ?? 0) > 0;
                const hasSingle = Boolean(params.agent && params.task);
                const modeCount = Number(hasTasks) + Number(hasSingle);

                const makeDetails =
                (mode: "single" | "parallel") =>
                (results: SingleResult[]): SubagentDetails => ({
                    mode,
                    results,
                });

                if (modeCount !== 1) {
                    const available = agents.map((a) => `${a.name} (${a.source})`).join(", ") || "none";
                    return {
                        content: [
                            {
                                type: "text",
                                text: `Invalid parameters. Provide exactly one mode.\nAvailable agents: ${available}`,
                            },
                        ],
                        details: makeDetails("single")([]),
                    };
                }

                if (params.tasks && params.tasks.length > 0) {
                    if (params.tasks.length > MAX_PARALLEL_TASKS)
                        return {
                            content: [
                                {
                                    type: "text",
                                    text: `Too many parallel tasks (${params.tasks.length}). Max is ${MAX_PARALLEL_TASKS}.`,
                                },
                            ],
                            details: makeDetails("parallel")([]),
                        };

                    const results = await runParallel(
                        ctx.cwd,
                        agents,
                        params.tasks,
                        ctx.modelRegistry,
                        signal,
                        (allResults) => {
                            if (!onUpdate) return;
                            const running = allResults.filter((r) => r.exitCode === -1).length;
                            const done = allResults.filter((r) => r.exitCode !== -1).length;
                            onUpdate({
                                content: [
                                    {
                                        type: "text",
                                        text: `Parallel: ${done}/${allResults.length} done, ${running} running...`,
                                    },
                                ],
                                details: makeDetails("parallel")(allResults),
                            });
                        },
                    );

                    const successCount = results.filter((r) => !isFailedResult(r)).length;
                    const summaries = results.map((r) => {
                        const output = truncateParallelOutput(getResultOutput(r));
                        const status = isFailedResult(r)
                            ? `failed${r.stopReason && r.stopReason !== "end" ? ` (${r.stopReason})` : ""}`
                            : "completed";
                        return `### [${r.agent}] ${status}\n\n${output}`;
                    });
                    return {
                        content: [
                            {
                                type: "text",
                                text: `Parallel: ${successCount}/${results.length} succeeded\n\n${summaries.join("\n\n---\n\n")}`,
                            },
                        ],
                        details: makeDetails("parallel")(results),
                    };
                }

                if (params.agent && params.task) {
                    const agent = agents.find((a) => a.name === params.agent);
                    if (!agent) {
                        const available = agents.map((a) => `"${a.name}"`).join(", ") || "none";
                        return {
                            content: [
                                {
                                    type: "text",
                                    text: `Unknown agent: "${params.agent}". Available agents: ${available}.`,
                                },
                            ],
                            details: makeDetails("single")([]),
                            isError: true,
                        };
                    }

                    const result = await runAgent(
                        agent,
                        params.task,
                        params.cwd ?? ctx.cwd,
                        ctx.modelRegistry,
                        signal,
                        onUpdate
                            ? (partial) => {
                                    onUpdate({
                                        content: [
                                            {
                                                type: "text",
                                                text: getFinalOutput(partial.messages) || "(running...)",
                                            },
                                        ],
                                        details: makeDetails("single")([partial]),
                                    });
                                }
                            : undefined,
                    );

                    const isError = isFailedResult(result);
                    if (isError) {
                        const errorMsg = getResultOutput(result);
                        return {
                            content: [{ type: "text", text: `Agent ${result.stopReason || "failed"}: ${errorMsg}` }],
                            details: makeDetails("single")([result]),
                            isError: true,
                        };
                    }
                    return {
                        content: [{ type: "text", text: getFinalOutput(result.messages) || "(no output)" }],
                        details: makeDetails("single")([result]),
                    };
                }

                const available = agents.map((a) => `${a.name} (${a.source})`).join(", ") || "none";
                return {
                    content: [{ type: "text", text: `Invalid parameters. Available agents: ${available}` }],
                    details: makeDetails("single")([]),
                };
            },

            renderCall(args, theme) {
                if (args.tasks && args.tasks.length > 0) {
                    let text =
                        theme.fg("toolTitle", theme.bold("delegate ")) +
                        theme.fg("accent", `parallel (${args.tasks.length} tasks)`);
                    for (const t of args.tasks.slice(0, 3)) {
                        const preview = t.task.length > 40 ? `${t.task.slice(0, 40)}...` : t.task;
                        text += `\n  ${theme.fg("accent", t.agent)}${theme.fg("dim", ` ${preview}`)}`;
                    }
                    if (args.tasks.length > 3)
                        text += `\n  ${theme.fg("muted", `... +${args.tasks.length - 3} more`)}`;
                    return new Text(text, 0, 0);
                }

                const agentName = args.agent || "...";
                const preview = args.task
                    ? args.task.length > 60
                        ? `${args.task.slice(0, 60)}...`
                        : args.task
                    : "...";
                const text =
                    theme.fg("toolTitle", theme.bold("delegate ")) +
                    theme.fg("accent", agentName) +
                    `\n  ${theme.fg("dim", preview)}`;
                return new Text(text, 0, 0);
            },

            renderResult(result, { expanded }, theme) {
                const details = result.details as SubagentDetails | undefined;
                if (!details || details.results.length === 0) {
                    const text = result.content[0];
                    return new Text(text?.type === "text" ? text.text : "(no output)", 0, 0);
                }

                const mdTheme = getMarkdownTheme();

                const renderDisplayItems = (items: DisplayItem[], limit?: number) => {
                    const toShow = limit ? items.slice(-limit) : items;
                    const skipped = limit && items.length > limit ? items.length - limit : 0;
                    let text = "";
                    if (skipped > 0) text += theme.fg("muted", `... ${skipped} earlier items\n`);
                    for (const item of toShow) {
                        if (item.type === "text") {
                            const preview = expanded ? item.text : item.text.split("\n").slice(0, 3).join("\n");
                            text += `${theme.fg("toolOutput", preview)}\n`;
                        } else {
                            text += `${theme.fg("muted", "→ ") + formatToolCall(item.name, item.args, theme.fg.bind(theme))}\n`;
                        }
                    }
                    return text.trimEnd();
                };

                if (details.mode === "single" && details.results.length === 1) {
                    const r = details.results[0];
                    const isError =
                        r.exitCode !== 0 || r.stopReason === "error" || r.stopReason === "aborted";
                    const icon = isError ? theme.fg("error", "✗") : theme.fg("success", "✓");
                    const displayItems = getDisplayItems(r.messages);
                    const finalOutput = getFinalOutput(r.messages);

                    if (expanded) {
                        const container = new Container();
                        let header = `${icon} ${theme.fg("toolTitle", theme.bold(r.agent))}${theme.fg("muted", ` (${r.agentSource})`)}`;
                        if (isError && r.stopReason) header += ` ${theme.fg("error", `[${r.stopReason}]`)}`;
                        container.addChild(new Text(header, 0, 0));
                        if (isError && r.errorMessage)
                            container.addChild(new Text(theme.fg("error", `Error: ${r.errorMessage}`), 0, 0));
                        container.addChild(new Spacer(1));
                        container.addChild(new Text(theme.fg("muted", "─── Task ───"), 0, 0));
                        container.addChild(new Text(theme.fg("dim", r.task), 0, 0));
                        container.addChild(new Spacer(1));
                        container.addChild(new Text(theme.fg("muted", "─── Output ───"), 0, 0));
                        if (displayItems.length === 0 && !finalOutput) {
                            container.addChild(new Text(theme.fg("muted", "(no output)"), 0, 0));
                        } else {
                            for (const item of displayItems) {
                                if (item.type === "toolCall")
                                    container.addChild(
                                        new Text(
                                            theme.fg("muted", "→ ") +
                                                formatToolCall(item.name, item.args, theme.fg.bind(theme)),
                                            0,
                                            0,
                                        ),
                                    );
                            }
                            if (finalOutput) {
                                container.addChild(new Spacer(1));
                                container.addChild(new Markdown(finalOutput.trim(), 0, 0, mdTheme));
                            }
                        }
                        const usageStr = formatUsageStats(r.usage, r.model);
                        if (usageStr) {
                            container.addChild(new Spacer(1));
                            container.addChild(new Text(theme.fg("dim", usageStr), 0, 0));
                        }
                        return container;
                    }

                    let text = `${icon} ${theme.fg("toolTitle", theme.bold(r.agent))}${theme.fg("muted", ` (${r.agentSource})`)}`;
                    if (isError && r.stopReason) text += ` ${theme.fg("error", `[${r.stopReason}]`)}`;
                    if (isError && r.errorMessage) text += `\n${theme.fg("error", `Error: ${r.errorMessage}`)}`;
                    else if (displayItems.length === 0) text += `\n${theme.fg("muted", "(no output)")}`;
                    else {
                        text += `\n${renderDisplayItems(displayItems, COLLAPSED_ITEM_COUNT)}`;
                        if (displayItems.length > COLLAPSED_ITEM_COUNT)
                            text += `\n${theme.fg("muted", "(Ctrl+O to expand)")}`;
                    }
                    const usageStr = formatUsageStats(r.usage, r.model);
                    if (usageStr) text += `\n${theme.fg("dim", usageStr)}`;
                    return new Text(text, 0, 0);
                }

                // Parallel mode
                const aggregateUsage = (results: SingleResult[]) => {
                    const total = { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, cost: 0, turns: 0 };
                    for (const r of results) {
                        total.input += r.usage.input;
                        total.output += r.usage.output;
                        total.cacheRead += r.usage.cacheRead;
                        total.cacheWrite += r.usage.cacheWrite;
                        total.cost += r.usage.cost;
                        total.turns += r.usage.turns;
                    }
                    return total;
                };

                const running = details.results.filter((r) => r.exitCode === -1).length;
                const successCount = details.results.filter((r) => r.exitCode === 0).length;
                const failCount = details.results.filter((r) => r.exitCode > 0).length;
                const isRunning = running > 0;
                const icon = isRunning
                    ? theme.fg("warning", "⏳")
                    : failCount > 0
                        ? theme.fg("warning", "◐")
                        : theme.fg("success", "✓");
                const status = isRunning
                    ? `${successCount + failCount}/${details.results.length} done, ${running} running`
                    : `${successCount}/${details.results.length} tasks`;

                if (expanded && !isRunning) {
                    const container = new Container();
                    container.addChild(
                        new Text(
                            `${icon} ${theme.fg("toolTitle", theme.bold("parallel "))}${theme.fg("accent", status)}`,
                            0,
                            0,
                        ),
                    );

                    for (const r of details.results) {
                        const rIcon = r.exitCode === 0 ? theme.fg("success", "✓") : theme.fg("error", "✗");
                        const displayItems = getDisplayItems(r.messages);
                        const finalOutput = getFinalOutput(r.messages);

                        container.addChild(new Spacer(1));
                        container.addChild(
                            new Text(`${theme.fg("muted", "─── ") + theme.fg("accent", r.agent)} ${rIcon}`, 0, 0),
                        );
                        container.addChild(
                            new Text(theme.fg("muted", "Task: ") + theme.fg("dim", r.task), 0, 0),
                        );
                        for (const item of displayItems) {
                            if (item.type === "toolCall")
                                container.addChild(
                                    new Text(
                                        theme.fg("muted", "→ ") +
                                            formatToolCall(item.name, item.args, theme.fg.bind(theme)),
                                        0,
                                        0,
                                    ),
                                );
                        }
                        if (finalOutput) {
                            container.addChild(new Spacer(1));
                            container.addChild(new Markdown(finalOutput.trim(), 0, 0, mdTheme));
                        }
                        const taskUsage = formatUsageStats(r.usage, r.model);
                        if (taskUsage) container.addChild(new Text(theme.fg("dim", taskUsage), 0, 0));
                    }

                    const usageStr = formatUsageStats(aggregateUsage(details.results));
                    if (usageStr) {
                        container.addChild(new Spacer(1));
                        container.addChild(new Text(theme.fg("dim", `Total: ${usageStr}`), 0, 0));
                    }
                    return container;
                }

                // Collapsed view (or still running)
                let text = `${icon} ${theme.fg("toolTitle", theme.bold("parallel "))}${theme.fg("accent", status)}`;
                for (const r of details.results) {
                    const rIcon =
                        r.exitCode === -1
                            ? theme.fg("warning", "⏳")
                            : r.exitCode === 0
                                ? theme.fg("success", "✓")
                                : theme.fg("error", "✗");
                    const displayItems = getDisplayItems(r.messages);
                    text += `\n\n${theme.fg("muted", "─── ")}${theme.fg("accent", r.agent)} ${rIcon}`;
                    if (displayItems.length === 0)
                        text += `\n${theme.fg("muted", r.exitCode === -1 ? "(running...)" : "(no output)")}`;
                    else text += `\n${renderDisplayItems(displayItems, 5)}`;
                }
                if (!isRunning) {
                    const usageStr = formatUsageStats(aggregateUsage(details.results));
                    if (usageStr) text += `\n\n${theme.fg("dim", `Total: ${usageStr}`)}`;
                }
                if (!expanded) text += `\n${theme.fg("muted", "(Ctrl+O to expand)")}`;
                return new Text(text, 0, 0);
            },
        });
    });
}