// Subagent Extension - Registers a tool for delegating work to prompt-defined
// subagents with constrained tool permissions.
import { spawn } from "node:child_process";
import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { fileURLToPath } from "node:url";
import type { ExtensionAPI, Theme } from "@mariozechner/pi-coding-agent";
import {
  parseFrontmatter,
  withFileMutationQueue,
} from "@mariozechner/pi-coding-agent";
import { Text } from "@mariozechner/pi-tui";
import { Type } from "typebox";

interface PromptConfig {
  name: string;
  description: string;
  approvedTools: string[];
  deniedTools: string[];
  systemPrompt: string;
  filePath: string;
}

interface SubagentToolActivity {
  id: string;
  toolName: string;
  summary: string;
  status: "running" | "done" | "error";
}

interface SubagentStatus {
  state: "starting" | "thinking" | "running" | "done";
  toolCallCount: number;
  activeToolCalls: number;
  recentToolCalls: SubagentToolActivity[];
}

interface SubagentResult {
  agent: string;
  task: string;
  tools: string[];
  exitCode: number;
  output: string;
  stderr: string;
  status: SubagentStatus;
  error?: string;
}

const EXTENSION_DIR = path.dirname(fileURLToPath(import.meta.url));
const PROMPTS_DIR = path.join(EXTENSION_DIR, "prompts");

// Parse Tool List - Frontmatter may use YAML arrays or comma-delimited strings.
function parseToolList(value: unknown): string[] {
  if (!value) return [];
  if (Array.isArray(value)) {
    return value
      .map(String)
      .map((tool) => tool.trim())
      .filter(Boolean);
  }
  if (typeof value === "string") {
    return value
      .split(",")
      .map((tool) => tool.trim())
      .filter(Boolean);
  }
  return [];
}

// Discover Prompts - Load prompt markdown files from this extension's prompts dir.
function discoverPrompts(): PromptConfig[] {
  if (!fs.existsSync(PROMPTS_DIR)) return [];

  const prompts: PromptConfig[] = [];
  const entries = fs.readdirSync(PROMPTS_DIR, { withFileTypes: true });
  for (const entry of entries) {
    if (!entry.name.endsWith(".md")) continue;
    if (!entry.isFile() && !entry.isSymbolicLink()) continue;

    const filePath = path.join(PROMPTS_DIR, entry.name);
    const content = fs.readFileSync(filePath, "utf-8");
    const { frontmatter, body } =
      parseFrontmatter<Record<string, unknown>>(content);

    if (
      typeof frontmatter.name !== "string" ||
      typeof frontmatter.description !== "string"
    ) {
      continue;
    }

    prompts.push({
      name: frontmatter.name,
      description: frontmatter.description,
      approvedTools: parseToolList(
        frontmatter.approved_tools ?? frontmatter.allowed_tools,
      ),
      deniedTools: parseToolList(frontmatter.denied_tools),
      systemPrompt: body.trim(),
      filePath,
    });
  }

  return prompts.sort((a, b) => a.name.localeCompare(b.name));
}

// Resolve Tools - Use exactly one permission mode. approved_tools/allowed_tools
// is a whitelist; denied_tools is a blacklist over the currently active tools.
function resolveTools(agent: PromptConfig, activeTools: string[]): string[] {
  if (agent.approvedTools.length > 0) {
    return [...new Set(agent.approvedTools)].filter(
      (tool) => tool !== "subagent",
    );
  }

  const denied = new Set(agent.deniedTools);
  return [...new Set(activeTools)].filter(
    (tool) => tool !== "subagent" && !denied.has(tool),
  );
}

// Write Prompt - pi accepts appended system prompts via file path.
async function writePromptToTempFile(
  agentName: string,
  prompt: string,
): Promise<{ dir: string; filePath: string }> {
  const dir = await fs.promises.mkdtemp(path.join(os.tmpdir(), "pi-subagent-"));
  const safeName = agentName.replace(/[^\w.-]+/g, "_");
  const filePath = path.join(dir, `${safeName}.md`);
  await withFileMutationQueue(filePath, async () => {
    await fs.promises.writeFile(filePath, prompt, {
      encoding: "utf-8",
      mode: 0o600,
    });
  });
  return { dir, filePath };
}

// Pi Invocation - Prefer the current pi entrypoint when running inside pi.
function getPiInvocation(args: string[]): { command: string; args: string[] } {
  const currentScript = process.argv[1];
  const isBunVirtualScript = currentScript?.startsWith("/$bunfs/root/");
  if (currentScript && !isBunVirtualScript && fs.existsSync(currentScript)) {
    return { command: process.execPath, args: [currentScript, ...args] };
  }

  const execName = path.basename(process.execPath).toLowerCase();
  const isGenericRuntime = /^(node|bun)(\.exe)?$/.test(execName);
  if (!isGenericRuntime) return { command: process.execPath, args };
  return { command: "pi", args };
}

// Truncate One Line - Keep live status compact and stable in the TUI.
function truncateOneLine(text: string, maxLength: number): string {
  const firstLine = text.split(/\r?\n/, 1)[0]?.trim() ?? "";
  return firstLine.length > maxLength
    ? `${firstLine.slice(0, Math.max(0, maxLength - 3))}...`
    : firstLine;
}

// Format Tool Args - Keep status updates compact enough for the tool UI.
function formatToolArgs(args: unknown): string {
  if (!args || typeof args !== "object") return "";

  const record = args as Record<string, unknown>;
  const priorityKeys = ["path", "command", "query", "pattern", "file", "agent"];
  const key = priorityKeys.find((candidate) => record[candidate] !== undefined);
  const value = key ? record[key] : undefined;
  if (value === undefined) return "";

  const text = typeof value === "string" ? value : JSON.stringify(value);
  const compact = text.replace(/\s+/g, " ").trim();
  return compact.length > 80 ? `${compact.slice(0, 77)}...` : compact;
}

// Format Status Text - Plain fallback for non-TUI consumers and JSON logs.
function formatStatusText(result: SubagentResult): string {
  const statusText =
    result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
  const lines = [
    `Subagent: ${result.agent}`,
    `Task: ${result.task}`,
    `Status: ${statusText}`,
    `Tool Calls: ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
  ];

  if (result.status.recentToolCalls.length > 0) {
    lines.push("", "Tool Calls - Last 3:");
    for (const call of result.status.recentToolCalls.slice(0, 3)) {
      const mark =
        call.status === "done" ? "✓" : call.status === "error" ? "✗" : "•";
      const suffix = call.summary ? `  ${call.summary}` : "";
      lines.push(`- ${mark} ${call.toolName}${suffix}`);
    }
  }

  return lines.join("\n");
}

// Render Subagent Summary - Compact collapsed view; Ctrl+O expands details.
function renderSubagentSummary(result: SubagentResult, theme: Theme): Text {
  const statusText =
    result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
  const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";
  const latest = result.status.recentToolCalls[0];
  const latestText = latest
    ? ` · ${latest.toolName}${latest.summary ? ` ${latest.summary}` : ""}`
    : "";
  const text =
    theme.fg(statusColor, statusText) +
    theme.fg(
      "muted",
      ` · ${result.status.toolCallCount} total, ${result.status.activeToolCalls} active${latestText}`,
    );
  return new Text(text, 0, 0);
}

// Render Subagent Status - Expanded TUI renderer with theme-aware colors.
function renderSubagentStatus(result: SubagentResult, theme: Theme): Text {
  const statusText =
    result.status.activeToolCalls > 0 ? "Running..." : "Thinking...";
  const statusColor = result.status.activeToolCalls > 0 ? "warning" : "muted";

  const lines = [
    `${theme.fg("dim", "Subagent:")} ${theme.fg("accent", result.agent)}`,
    `${theme.fg("dim", "Task:")} ${theme.fg("muted", result.task)}`,
    `${theme.fg("dim", "Status:")} ${theme.fg(statusColor, statusText)}`,
    `${theme.fg("dim", "Tool Calls:")} ${theme.fg(
      "muted",
      `${result.status.toolCallCount} total, ${result.status.activeToolCalls} active`,
    )}`,
  ];

  if (result.status.recentToolCalls.length > 0) {
    lines.push("", theme.fg("dim", "Tool Calls - Last 3:"));
    for (const call of result.status.recentToolCalls.slice(0, 3)) {
      const mark =
        call.status === "done"
          ? theme.fg("success", "✓")
          : call.status === "error"
            ? theme.fg("error", "✗")
            : theme.fg("warning", "•");
      const tool = theme.fg("toolTitle", call.toolName.padEnd(6));
      const summary = call.summary ? ` ${theme.fg("muted", call.summary)}` : "";
      lines.push(`- ${mark} ${tool}${summary}`);
    }
  }

  return new Text(lines.join("\n"), 0, 0);
}

// Run Agent - Spawn an isolated pi process in JSON mode and collect final text.
async function runAgent(
  cwd: string,
  agent: PromptConfig,
  task: string,
  tools: string[],
  signal?: AbortSignal,
  onUpdate?: (partial: {
    content: { type: "text"; text: string }[];
    details: SubagentResult;
  }) => void,
): Promise<SubagentResult> {
  let tmpDir: string | null = null;
  let tmpPromptPath: string | null = null;

  const result: SubagentResult = {
    agent: agent.name,
    task,
    tools,
    exitCode: 0,
    output: "",
    stderr: "",
    status: {
      state: "starting",
      toolCallCount: 0,
      activeToolCalls: 0,
      recentToolCalls: [],
    },
  };

  const activeToolIds = new Set<string>();

  const emitUpdate = () => {
    onUpdate?.({
      content: [{ type: "text", text: formatStatusText(result) }],
      details: {
        ...result,
        status: {
          ...result.status,
          recentToolCalls: [...result.status.recentToolCalls],
        },
      },
    });
  };

  const rememberToolCall = (activity: SubagentToolActivity) => {
    const existing = result.status.recentToolCalls.findIndex(
      (call) => call.id === activity.id,
    );
    if (existing >= 0) result.status.recentToolCalls.splice(existing, 1);
    result.status.recentToolCalls.unshift(activity);
    result.status.recentToolCalls = result.status.recentToolCalls.slice(0, 3);
  };

  try {
    const args = ["--mode", "json", "-p", "--no-session"];
    args.push("--tools", tools.join(","));

    if (agent.systemPrompt) {
      const tmp = await writePromptToTempFile(agent.name, agent.systemPrompt);
      tmpDir = tmp.dir;
      tmpPromptPath = tmp.filePath;
      args.push("--append-system-prompt", tmpPromptPath);
    }

    args.push(`Task: ${task}`);
    emitUpdate();

    const exitCode = await new Promise<number>((resolve) => {
      const invocation = getPiInvocation(args);
      const proc = spawn(invocation.command, invocation.args, {
        cwd,
        shell: false,
        stdio: ["ignore", "pipe", "pipe"],
      });

      let buffer = "";
      let aborted = false;

      const processLine = (line: string) => {
        if (!line.trim()) return;
        try {
          const event = JSON.parse(line);
          if (
            event.type === "message_start" &&
            event.message?.role === "assistant"
          ) {
            result.status.state = "thinking";
            emitUpdate();
          }
          if (event.type === "tool_execution_start") {
            const id = String(
              event.toolCallId ?? result.status.toolCallCount + 1,
            );
            activeToolIds.add(id);
            result.status.state = "running";
            result.status.toolCallCount += 1;
            result.status.activeToolCalls = activeToolIds.size;
            rememberToolCall({
              id,
              toolName: String(event.toolName ?? "tool"),
              summary: formatToolArgs(event.args),
              status: "running",
            });
            emitUpdate();
          }
          if (event.type === "tool_execution_end") {
            const id = String(event.toolCallId ?? "");
            if (id) activeToolIds.delete(id);
            result.status.activeToolCalls = activeToolIds.size;
            const previous = result.status.recentToolCalls.find(
              (call) => call.id === id,
            );
            rememberToolCall({
              id,
              toolName: String(event.toolName ?? previous?.toolName ?? "tool"),
              summary: previous?.summary ?? "",
              status: event.isError ? "error" : "done",
            });
            emitUpdate();
          }
          if (
            event.type === "message_end" &&
            event.message?.role === "assistant"
          ) {
            for (const part of event.message.content ?? []) {
              if (part.type === "text") result.output = part.text;
            }
            if (event.message.stopReason === "error") {
              result.error =
                event.message.errorMessage ?? "Subagent stopped with an error.";
            }
            emitUpdate();
          }
        } catch {
          // Ignore non-JSON output.
        }
      };

      proc.stdout.on("data", (data) => {
        buffer += data.toString();
        const lines = buffer.split("\n");
        buffer = lines.pop() ?? "";
        for (const line of lines) processLine(line);
      });

      proc.stderr.on("data", (data) => {
        result.stderr += data.toString();
      });

      proc.on("close", (code) => {
        if (buffer.trim()) processLine(buffer);
        resolve(aborted ? 130 : (code ?? 0));
      });

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

      if (signal) {
        const killProc = () => {
          aborted = true;
          proc.kill("SIGTERM");
          setTimeout(() => {
            if (!proc.killed) proc.kill("SIGKILL");
          }, 5000);
        };
        if (signal.aborted) killProc();
        else signal.addEventListener("abort", killProc, { once: true });
      }
    });

    result.exitCode = exitCode;
    result.status.state = "done";
    result.status.activeToolCalls = 0;
    if (exitCode === 130) result.error = "Subagent was aborted.";
    return result;
  } finally {
    if (tmpPromptPath) await fs.promises.rm(tmpPromptPath, { force: true });
    if (tmpDir) await fs.promises.rm(tmpDir, { force: true, recursive: true });
  }
}

// Format Prompt List - Keep available prompt hints compact for the LLM.
function formatPromptList(prompts: PromptConfig[]): string {
  if (prompts.length === 0) return "none";
  return prompts
    .map((prompt) => `${prompt.name}: ${prompt.description}`)
    .join("; ");
}

const SubagentParams = Type.Object({
  agent: Type.String({
    description: "Name of the prompt-defined subagent to invoke",
  }),
  task: Type.String({ description: "Task to delegate to the subagent" }),
  cwd: Type.Optional(
    Type.String({
      description:
        "Working directory for the subagent process. Defaults to current cwd.",
    }),
  ),
});

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "subagent",
    label: "Subagent",
    description:
      "Delegate a task to a prompt-defined subagent from this extension's prompts/ directory. " +
      `Available at startup: ${formatPromptList(discoverPrompts())}`,
    promptSnippet:
      "Delegate tasks to subagents by name. Subagent prompts live in prompts/*.md and define approved_tools/denied_tools.",
    parameters: SubagentParams,

    async execute(_toolCallId, params, signal, _onUpdate, ctx) {
      const prompts = discoverPrompts();
      const agent = prompts.find((prompt) => prompt.name === params.agent);
      if (!agent) {
        return {
          content: [
            {
              type: "text",
              text: `Unknown subagent: ${params.agent}. Available: ${formatPromptList(prompts)}`,
            },
          ],
          details: { available: prompts.map((prompt) => prompt.name) },
          isError: true,
        };
      }

      if (agent.approvedTools.length > 0 && agent.deniedTools.length > 0) {
        return {
          content: [
            {
              type: "text",
              text:
                `Invalid subagent config for ${agent.name}: define either approved_tools/allowed_tools ` +
                "or denied_tools, not both.",
            },
          ],
          details: { agent: agent.name, filePath: agent.filePath },
          isError: true,
        };
      }

      const activeTools = pi.getActiveTools();
      const tools = resolveTools(agent, activeTools);
      if (tools.length === 0) {
        return {
          content: [
            {
              type: "text",
              text: `Subagent ${agent.name} has no approved tools after applying denied_tools.`,
            },
          ],
          details: {
            agent: agent.name,
            approvedTools: agent.approvedTools,
            deniedTools: agent.deniedTools,
          },
          isError: true,
        };
      }

      const cwd = path.resolve(ctx.cwd, params.cwd ?? ".");
      const result = await runAgent(
        cwd,
        agent,
        params.task,
        tools,
        signal,
        _onUpdate,
      );
      const failed = result.exitCode !== 0 || Boolean(result.error);
      const fallback = result.error || result.stderr || "(no output)";

      return {
        content: [{ type: "text", text: result.output || fallback }],
        details: result,
        isError: failed,
      };
    },

    renderCall(args, theme) {
      const task = truncateOneLine(args.task, 90);
      const text =
        theme.fg("toolTitle", theme.bold("subagent ")) +
        theme.fg("accent", args.agent) +
        (task ? ` ${theme.fg("dim", task)}` : "");
      return new Text(text, 0, 0);
    },

    renderResult(result, { expanded, isPartial }, theme) {
      const details = result.details as Partial<SubagentResult> | undefined;
      if (isPartial && details?.status && details.agent && details.task) {
        return expanded
          ? renderSubagentStatus(details as SubagentResult, theme)
          : renderSubagentSummary(details as SubagentResult, theme);
      }

      const text = result.content[0];
      return new Text(text?.type === "text" ? text.text : "", 0, 0);
    },
  });
}