import { spawn } from "node:child_process";
import { Type } from "@sinclair/typebox";
import { StringEnum } from "@mariozechner/pi-ai";
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { trimOutput } from "@hyperprior/pi-shared";
import {
  makeSubagentResultDetails,
  normalizeSubagentRunResult,
  type SubagentMode,
  type SubagentResultDetails,
  type SubagentRunResult,
} from "../src/result-schema.ts";

const MAX_CONCURRENCY = 4;
const MAX_PARALLEL_TASKS = 12;
const DEFAULT_TIMEOUT_MS = 180_000;

const SubagentTask = Type.Object(
  {
    task: Type.String({ description: "Prompt/task for one subagent run" }),
    model: Type.Optional(Type.String({ description: "Model override for this run" })),
    tools: Type.Optional(Type.String({ description: "Comma-separated tools for this run" })),
    cwd: Type.Optional(Type.String({ description: "Working directory for this run" })),
    agent: Type.Optional(Type.String({ description: "Agent label for display only" })),
  },
  { additionalProperties: false },
);

const SubagentParams = Type.Object({
  mode: StringEnum(["single", "parallel", "chain"] as const, { description: "Execution mode" }),
  task: Type.Optional(Type.String({ description: "Prompt for single mode" })),
  tasks: Type.Optional(Type.Array(SubagentTask, { description: "Tasks for parallel mode" })),
  chain: Type.Optional(Type.Array(SubagentTask, { description: "Tasks for chain mode" })),
  model: Type.Optional(Type.String({ description: "Model used for all runs unless overridden per task" })),
  tools: Type.Optional(Type.String({ description: "Comma-separated tools for all runs unless overridden per task" })),
  cwd: Type.Optional(Type.String({ description: "Default working directory for all runs" })),
  timeout_ms: Type.Optional(
    Type.Number({ minimum: 1000, maximum: 900_000, default: DEFAULT_TIMEOUT_MS, description: "Kill timeout per run (ms)" }),
  ),
});

type SubagentTaskInput = {
  task: string;
  model?: string;
  tools?: string;
  cwd?: string;
  agent?: string;
};

interface NormalizedTask extends SubagentTaskInput {
  id: number;
  finalTask: string;
  effectiveModel?: string;
  effectiveTools?: string;
  effectiveCwd: string;
}

function normalizeTask(task: SubagentTaskInput, index: number, defaults: {
  model?: string;
  tools?: string;
  cwd?: string;
}, previous?: string): NormalizedTask {
  const raw = previous ? task.task.replace(/\{previous\}/g, previous) : task.task;
  return {
    id: index + 1,
    task: raw,
    finalTask: raw,
    agent: task.agent?.trim() || `agent-${index + 1}`,
    effectiveModel: (task.model || defaults.model)?.trim() || undefined,
    effectiveTools: (task.tools || defaults.tools)?.trim() || undefined,
    effectiveCwd: task.cwd?.trim() || defaults.cwd?.trim() || process.cwd(),
  };
}

function toCliResult(raw: any): string | null {
  if (!raw || typeof raw !== "object") return null;
  const payload = raw.message ?? raw;
  if (!payload || payload.role !== "assistant") return null;

  const content = Array.isArray(payload.content) ? payload.content : [];
  for (const part of content) {
    if (part?.type === "text" && typeof part.text === "string") {
      return part.text;
    }
  }
  return null;
}

async function runSubagentPrompt(task: NormalizedTask, timeoutMs: number, signal: AbortSignal | undefined): Promise<SubagentRunResult> {
  const start = Date.now();
  const args: string[] = ["-p", "--no-session", "--mode", "json"];

  if (task.effectiveModel) args.push("--model", task.effectiveModel);
  if (task.effectiveTools) args.push("--tools", task.effectiveTools);

  args.push(task.finalTask);

  let stdout = "";
  let stderr = "";
  let finalOutput = "";
  let finalUsage = "";
  let status = "running";

  const child = spawn("pi", args, {
    cwd: task.effectiveCwd,
    stdio: ["ignore", "pipe", "pipe"],
    shell: false,
  });

  const timer = setTimeout(() => {
    child.kill("SIGTERM");
  }, timeoutMs);

  if (signal) {
    if (signal.aborted) child.kill("SIGTERM");
    else {
      signal.addEventListener(
        "abort",
        () => {
          child.kill("SIGTERM");
        },
        { once: true },
      );
    }
  }

  let stdoutBuffer = "";

  child.stdout.setEncoding("utf8");
  child.stderr.setEncoding("utf8");

  const consumeLine = (line: string) => {
    if (!line.trim()) return;
    let event: any;
    try {
      event = JSON.parse(line);
    } catch {
      return;
    }
    const value = toCliResult(event);
    if (value) finalOutput = value;
    if (typeof event?.finalUsage === "string") finalUsage = event.finalUsage;
    const statusEvent = event?.type === "message_end" ? (event.message ?? event) : null;
    if (statusEvent && typeof statusEvent === "object" && (statusEvent as any).stopReason) {
      status = (statusEvent as any).stopReason;
    }
  };

  child.stdout.on("data", (chunk) => {
    stdout += chunk;
    stdoutBuffer += chunk;
    const lines = stdoutBuffer.split("\n");
    stdoutBuffer = lines.pop() || "";
    for (const line of lines) {
      consumeLine(line);
    }
  });

  child.stderr.on("data", (chunk) => {
    stderr += chunk;
  });

  const exitCode = await new Promise<number>((resolve) => {
    child.on("close", (code) => resolve(code ?? 0));
    child.on("error", () => resolve(1));
  });

  clearTimeout(timer);
  if (stdoutBuffer.trim()) {
    consumeLine(stdoutBuffer);
  }

  const elapsedMs = Date.now() - start;
  const normalized = trimOutput(finalOutput || finalUsage || stdout || stderr || "", 12_000);

  if (exitCode !== 0) {
    return normalizeSubagentRunResult({
      task: task.finalTask,
      agent: task.agent || "subagent",
      exitCode,
      elapsedMs,
      output: normalized,
      error: stderr || normalized,
      model: task.effectiveModel,
      tools: task.effectiveTools,
    });
  }

  if (status === "error" || status === "aborted") {
    return normalizeSubagentRunResult({
      task: task.finalTask,
      agent: task.agent || "subagent",
      exitCode: exitCode || 1,
      elapsedMs,
      output: normalized,
      error: status,
      model: task.effectiveModel,
      tools: task.effectiveTools,
    });
  }

  return normalizeSubagentRunResult({
    task: task.finalTask,
    agent: task.agent || "subagent",
    exitCode,
    elapsedMs,
    output: normalized,
    model: task.effectiveModel,
    tools: task.effectiveTools,
  });
}

function formatRunSummary(result: SubagentRunResult & { taskIndex: number }): string {
  const status = result.exitCode === 0 ? "✓" : "✗";
  const location = result.agent ? ` [${result.agent}]` : "";
  const model = result.model ? ` (${result.model})` : "";
  const head = result.output.split("\n").slice(0, 2).join(" ");
  return `${status} #${result.taskIndex}${location}${model}: ${trimOutput(head, 180)}${result.exitCode === 0 ? "" : `\n  ⚠ ${result.error ?? "failed"}`}`;
}

async function runWithConcurrency<T, R>(
  items: T[],
  concurrency: number,
  handler: (item: T, index: number) => Promise<R>,
): Promise<R[]> {
  const limit = Math.max(1, Math.min(concurrency, items.length));
  const results = new Array<R>(items.length);
  let next = 0;

  await Promise.all(
    Array.from({ length: limit }, async () => {
      while (true) {
        const current = next++;
        if (current >= items.length) return;
        results[current] = await handler(items[current], current);
      }
    }),
  );

  return results;
}

export default function (pi: ExtensionAPI) {
  pi.registerTool({
    name: "hyperpi_subagent",
    label: "Subagent",
    description:
      "Run isolated pi sessions in subprocesses. Modes: single, parallel, chain. " +
      "Each subtask gets a full model/tool context and returns structured output.",
    parameters: SubagentParams,
    async execute(_toolCallId, params, signal, onUpdate, ctx) {
      const defaults = {
        model: params.model?.trim(),
        tools: params.tools?.trim(),
        cwd: params.cwd?.trim() || ctx.cwd,
      };
      const timeoutMs = params.timeout_ms || DEFAULT_TIMEOUT_MS;
      const mode = params.mode;

      const emptyDetails: SubagentResultDetails = makeSubagentResultDetails(mode, []);

      const emit = (tasks: SubagentRunResult[]) => {
        if (!onUpdate) return;
        onUpdate({
          content: [
            {
              type: "text",
              text: `${mode}: ${tasks.filter((t) => t.exitCode === 0).length}/${tasks.length} completed`,
            },
          ],
          details: makeSubagentResultDetails(mode, tasks),
        });
      };

      if (mode === "single") {
        if (!params.task?.trim()) {
          return {
            content: [{ type: "text", text: "Single mode requires `task`." }],
            details: emptyDetails,
            isError: true,
          };
        }

        const task = normalizeTask(
          {
            task: params.task,
            model: defaults.model,
            tools: defaults.tools,
            cwd: defaults.cwd,
          },
          0,
          defaults,
        );
        const result = await runSubagentPrompt(task, timeoutMs, signal);
        return {
          content: [{ type: "text", text: result.output || "(no output)" }],
          details: makeSubagentResultDetails(mode, [result]),
          isError: result.exitCode !== 0,
        };
      }

      if (mode === "parallel") {
        if (!params.tasks?.length) {
          return {
            content: [{ type: "text", text: "Parallel mode requires `tasks` array." }],
            details: emptyDetails,
            isError: true,
          };
        }

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

        const normalized = params.tasks.map((item, idx) => normalizeTask(item, idx, defaults));
        const running: SubagentRunResult[] = normalized.map((task) =>
          normalizeSubagentRunResult({
            task: task.finalTask,
            exitCode: -1,
            agent: task.agent || `agent-${task.id}`,
            output: "(queued)",
            elapsedMs: 0,
            model: task.effectiveModel,
            tools: task.effectiveTools,
          }),
        );

        const update = () => {
          if (onUpdate) {
            emit(running);
          }
        };

        const results = await runWithConcurrency(normalized, MAX_CONCURRENCY, async (task, index) => {
          const result = await runSubagentPrompt(task, timeoutMs, signal);
          running[index] = result;
          update();
          return result;
        });

        return {
          content: [{
            type: "text",
            text: results
              .map((result, i) => formatRunSummary({ ...result, taskIndex: i + 1 }))
              .join("\n"),
          }],
          details: makeSubagentResultDetails(mode, results),
          isError: results.some((result) => result.exitCode !== 0),
        };
      }

      if (mode === "chain") {
        if (!params.chain?.length) {
          return {
            content: [{ type: "text", text: "Chain mode requires `chain` array." }],
            details: emptyDetails,
            isError: true,
          };
        }

        const results: Array<SubagentRunResult> = [];
        let previous = "";

        for (let i = 0; i < params.chain.length; i++) {
          const task = normalizeTask(params.chain[i]!, i, defaults, previous);
          const output = await runSubagentPrompt(task, timeoutMs, signal);
          results.push(output);

          const payload = `#${i + 1} ${task.agent || "agent"}: ${output.exitCode === 0 ? "ok" : "failed"}`;
          onUpdate?.({
            content: [{ type: "text", text: payload }],
            details: makeSubagentResultDetails(mode, results),
          });

          if (output.exitCode !== 0) {
            return {
              content: [{
                type: "text",
                text: `Chain halted at step ${i + 1}: ${output.error || trimOutput(output.output, 200)}`,
              }],
              details: makeSubagentResultDetails(mode, results),
              isError: true,
            };
          }

          previous = output.output.slice(0, 160);
        }

        return {
          content: [{
            type: "text",
            text: results
              .map((result, i) => formatRunSummary({ ...result, taskIndex: i + 1 }))
              .join("\n\n"),
          }],
          details: makeSubagentResultDetails(mode, results),
        };
      }

      return {
        content: [{ type: "text", text: `Unsupported mode: ${mode}` }],
        details: emptyDetails,
        isError: true,
      };
    },
  });

  pi.registerCommand("subagent", {
    description: "Run single-line subagent task in a separate session (single mode).",
    handler: async (args, ctx) => {
      const task = (args || "").trim();
      if (!task) {
        if (ctx.hasUI) {
          ctx.ui.notify("Usage: /subagent <task>", "warning");
        }
        return;
      }

      const result = await runSubagentPrompt(
        {
          id: 1,
          task,
          agent: "inline",
          finalTask: task,
          effectiveCwd: process.cwd(),
        },
        DEFAULT_TIMEOUT_MS,
        undefined,
      );

      if (ctx.hasUI) {
        ctx.ui.notify(result.output || "(no output)", result.exitCode === 0 ? "info" : "error");
      }
      return result.output || "(no output)";
    },
  });
}