// Subagent Extension - Registers a tool for delegating work to prompt-defined
// subagents with constrained tool permissions.
import { spawn } from "node:child_process";
import { createHash, randomUUID } from "node:crypto";
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 { Markdown, Text, type MarkdownTheme } 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 SubagentFinalizePayload {
  status: "SUCCESS" | "ERROR";
  result?: string;
  error?: string;
}

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

const EXTENSION_ENTRY = fileURLToPath(import.meta.url);
const EXTENSION_DIR = path.dirname(EXTENSION_ENTRY);
const PROMPTS_DIR = path.join(EXTENSION_DIR, "prompts");
const FINALIZE_TOOL_NAME = "subagent_finalize";
const MAX_FINALIZE_RETRIES = 2;

// Markdown Theme - Render finalized subagent text with Pi's terminal markdown component.
function getSubagentMarkdownTheme(theme: Theme): MarkdownTheme {
  return {
    heading: (text) => theme.fg("mdHeading", text),
    link: (text) => theme.fg("mdLink", text),
    linkUrl: (text) => theme.fg("mdLinkUrl", text),
    code: (text) => theme.fg("mdCode", text),
    codeBlock: (text) => theme.fg("mdCodeBlock", text),
    codeBlockBorder: (text) => theme.fg("mdCodeBlockBorder", text),
    quote: (text) => theme.fg("mdQuote", text),
    quoteBorder: (text) => theme.fg("mdQuoteBorder", text),
    hr: (text) => theme.fg("mdHr", text),
    listBullet: (text) => theme.fg("mdListBullet", text),
    bold: (text) => theme.bold(text),
    italic: (text) => theme.italic(text),
    strikethrough: (text) => theme.strikethrough(text),
    underline: (text) => theme.underline(text),
    codeBlockIndent: "  ",
  };
}

// Format Tool Content - Some clients hide structured details from the model.
function formatSubagentContent(
  status: "SUCCESS" | "ERROR",
  sessionId: string,
  result?: string,
  error?: string,
): string {
  const header = [`**Status:** ${status}`, `**Session ID:** \`${sessionId}\``];
  if (error?.trim()) header.push(`**Error:** ${error.trim()}`);

  const body = result?.trimEnd();
  return body
    ? `${header.join("  \n")}\n\n---\n\n${body}`
    : header.join("  \n");
}

// 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[] {
  const withoutDelegationTools = (tools: string[]) =>
    tools.filter((tool) => tool !== "subagent" && tool !== FINALIZE_TOOL_NAME);

  const resolved =
    agent.approvedTools.length > 0
      ? withoutDelegationTools([...new Set(agent.approvedTools)])
      : withoutDelegationTools(
          [...new Set(activeTools)].filter(
            (tool) => !new Set(agent.deniedTools).has(tool),
          ),
        );

  return [...new Set([...resolved, FINALIZE_TOOL_NAME])];
}

// Build Finalize Prompt - Child agents must terminate by calling this tool.
function buildSubagentPrompt(agent: PromptConfig): string {
  const finalizePrompt = [
    "You are running as a subagent.",
    `When the task is complete, call ${FINALIZE_TOOL_NAME} as your final action.`,
    "Do not provide the final answer as normal assistant text.",
    `${FINALIZE_TOOL_NAME} requires status SUCCESS with result, or status ERROR with error and optional result.`,
  ].join("\n");

  return [agent.systemPrompt, finalizePrompt].filter(Boolean).join("\n\n");
}

// Session Path - Persist child sessions as <agent>_<uuid>.jsonl under a cwd hash.
function getSubagentSessionPath(
  cwd: string,
  agentName: string,
  sessionId: string,
): string {
  const cwdHash = createHash("sha256").update(cwd).digest("hex").slice(0, 16);
  const safeAgent = agentName.replace(/[^\w.-]+/g, "_");
  const safeSessionId = sessionId.replace(/[^\w.-]+/g, "_");
  return path.join(
    os.homedir(),
    ".pi",
    "subagent-sessions",
    cwdHash,
    `${safeAgent}_${safeSessionId}.jsonl`,
  );
}

// Validate Finalize Payload - Keep the parent contract strict and small.
function validateFinalizePayload(
  value: unknown,
): SubagentFinalizePayload | null {
  if (!value || typeof value !== "object") return null;
  const payload = value as Record<string, unknown>;
  if (payload.status === "SUCCESS") {
    return typeof payload.result === "string" && payload.result.trim()
      ? { status: "SUCCESS", result: payload.result }
      : null;
  }
  if (payload.status === "ERROR") {
    const result =
      typeof payload.result === "string" ? payload.result : undefined;
    return typeof payload.error === "string" && payload.error.trim()
      ? { status: "ERROR", error: payload.error, result }
      : null;
  }
  return null;
}

// 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[],
  sessionId: string,
  sessionPath: 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,
    sessionId,
    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",
      "--session",
      sessionPath,
      "--extension",
      EXTENSION_ENTRY,
    ];
    args.push("--tools", tools.join(","));

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

    args.push(task);
    emitUpdate();

    const exitCode = await new Promise<number>((resolve) => {
      const invocation = getPiInvocation(args);
      const proc = spawn(invocation.command, invocation.args, {
        cwd,
        env: { ...process.env, PI_SUBAGENT_CHILD: "1" },
        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,
            );
            const toolName = String(event.toolName ?? "tool");
            if (toolName === FINALIZE_TOOL_NAME) {
              result.finalized =
                validateFinalizePayload(event.args) ?? undefined;
            }
            activeToolIds.add(id);
            result.status.state = "running";
            result.status.toolCallCount += 1;
            result.status.activeToolCalls = activeToolIds.size;
            rememberToolCall({
              id,
              toolName,
              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" }),
  sessionId: Type.Optional(
    Type.String({
      description:
        "Optional sticky subagent session id. Reuse to continue a previous subagent context.",
    }),
  ),
  cwd: Type.Optional(
    Type.String({
      description:
        "Working directory for the subagent process. Defaults to current cwd.",
    }),
  ),
});

export default function (pi: ExtensionAPI) {
  if (process.env.PI_SUBAGENT_CHILD === "1") {
    pi.registerTool({
      name: FINALIZE_TOOL_NAME,
      label: "Subagent Finalize",
      description:
        "Internal subagent-only tool. Call this as your final action when delegated work is complete.",
      promptSnippet:
        "Call subagent_finalize as your final action when subagent work is complete.",
      parameters: Type.Object({
        status: Type.Union([Type.Literal("SUCCESS"), Type.Literal("ERROR")]),
        result: Type.Optional(Type.String()),
        error: Type.Optional(Type.String()),
      }),

      async execute(_toolCallId, params) {
        return {
          content: [{ type: "text", text: "Subagent finalized." }],
          details: params,
          terminate: true,
        };
      },
    });
  }

  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 sessionId = params.sessionId ?? randomUUID();
      const sessionPath = getSubagentSessionPath(cwd, agent.name, sessionId);
      await fs.promises.mkdir(path.dirname(sessionPath), { recursive: true });

      let result: SubagentResult | null = null;
      for (
        let retryCount = 0;
        retryCount <= MAX_FINALIZE_RETRIES;
        retryCount += 1
      ) {
        const task =
          retryCount === 0
            ? `Task: ${params.task}`
            : [
                "Your previous response did not finalize correctly.",
                `If you are finished, call ${FINALIZE_TOOL_NAME}.`,
                "If you are not finished, continue the original task using available tools as needed.",
                `Original task: ${params.task}`,
              ].join("\n\n");

        result = await runAgent(
          cwd,
          agent,
          task,
          tools,
          sessionId,
          sessionPath,
          signal,
          _onUpdate,
        );

        if (result.finalized) break;
        if (result.exitCode !== 0 || result.error) break;
      }

      if (!result) {
        return {
          content: [
            {
              type: "text",
              text: formatSubagentContent(
                "ERROR",
                sessionId,
                undefined,
                "Subagent did not run.",
              ),
            },
          ],
          details: { sessionId },
          isError: true,
        };
      }

      if (!result.finalized) {
        const fallback =
          result.error ||
          result.stderr ||
          `Subagent did not call ${FINALIZE_TOOL_NAME}.`;
        return {
          content: [
            {
              type: "text",
              text: formatSubagentContent(
                "ERROR",
                sessionId,
                undefined,
                fallback,
              ),
            },
          ],
          details: { sessionId },
          isError: true,
        };
      }

      if (result.finalized.status === "ERROR") {
        return {
          content: [
            {
              type: "text",
              text: formatSubagentContent(
                "ERROR",
                sessionId,
                result.finalized.result,
                result.finalized.error ?? "Subagent failed.",
              ),
            },
          ],
          details: { sessionId, finalized: result.finalized },
          isError: true,
        };
      }

      return {
        content: [
          {
            type: "text",
            text: formatSubagentContent(
              "SUCCESS",
              sessionId,
              result.finalized.result,
            ),
          },
        ],
        details: { sessionId, finalized: result.finalized },
        isError: false,
      };
    },

    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 Markdown(
        text?.type === "text" ? text.text : "",
        0,
        0,
        getSubagentMarkdownTheme(theme),
      );
    },
  });
}