import type { ExtensionAPI } from "@jaeswift/jae-coding-agent";
import { spawn, type ChildProcess } from "node:child_process";
import { Type } from "@sinclair/typebox";

interface AgentInstance {
  id: string;
  role: string;
  systemPrompt: string;
  status: "running" | "idle" | "stopped" | "error";
  process: ChildProcess | null;
  output: string;
  createdAt: number;
  lastMessage: string;
}

const ROLE_PROMPTS: Record<string, string> = {
  frontend: "You are a frontend development specialist. Focus on UI/UX, React, CSS, accessibility, responsive design, and browser APIs. Write clean component code with proper state management.",
  backend: "You are a backend development specialist. Focus on API design, databases, server architecture, authentication, and performance. Write robust server-side code with proper error handling.",
  tester: "You are a QA and testing specialist. Focus on writing comprehensive tests: unit tests, integration tests, E2E tests. Identify edge cases, write test plans, and ensure code reliability.",
  reviewer: "You are a senior code reviewer. Analyze code for bugs, security issues, performance problems, and style violations. Provide constructive feedback with specific suggestions.",
  devops: "You are a DevOps and infrastructure specialist. Focus on CI/CD pipelines, Docker, Kubernetes, monitoring, and deployment automation. Write infrastructure as code.",
  security: "You are a security specialist. Focus on vulnerability assessment, secure coding practices, authentication/authorization, encryption, and threat modeling.",
  architect: "You are a software architect. Focus on system design, scalability patterns, technology selection, API contracts, and architectural decision records.",
  docs: "You are a technical documentation specialist. Focus on writing clear README files, API docs, tutorials, inline comments, and architectural documentation.",
};

export default function (pi: ExtensionAPI) {
  const agents = new Map<string, AgentInstance>();
  let idCounter = 0;

  function generateId(role: string): string {
    idCounter++;
    return `${role}-${idCounter}`;
  }

  function getSystemPrompt(role: string): string {
    return ROLE_PROMPTS[role] || `You are a ${role} specialist. Focus on tasks related to ${role} and provide expert-level assistance.`;
  }

  async function spawnAgent(role: string, cwd: string): Promise<AgentInstance> {
    const id = generateId(role);
    const systemPrompt = getSystemPrompt(role);

    const agent: AgentInstance = {
      id,
      role,
      systemPrompt,
      status: "idle",
      process: null,
      output: "",
      createdAt: Date.now(),
      lastMessage: "",
    };

    agents.set(id, agent);
    return agent;
  }

  async function sendMessage(
    agent: AgentInstance,
    message: string,
    cwd: string,
  ): Promise<string> {
    return new Promise((resolve) => {
      const timeout = 180_000; // 3 minutes
      let output = "";
      let resolved = false;

      agent.status = "running";
      agent.lastMessage = message;

      const prompt = `[System: ${agent.systemPrompt}]\n\n${message}`;

      try {
        const child = spawn(
          "jae",
          ["-p", prompt],
          {
            cwd,
            stdio: ["ignore", "pipe", "pipe"],
            timeout,
            env: { ...process.env },
          },
        );

        agent.process = child;

        child.stdout?.on("data", (data: Buffer) => {
          output += data.toString();
        });

        child.stderr?.on("data", (data: Buffer) => {
          output += data.toString();
        });

        child.on("close", (code) => {
          agent.process = null;
          agent.status = code === 0 ? "idle" : "error";
          agent.output = output.trim();
          if (!resolved) {
            resolved = true;
            resolve(output.trim() || "(no output)");
          }
        });

        child.on("error", (err) => {
          agent.process = null;
          agent.status = "error";
          if (!resolved) {
            resolved = true;
            resolve(`Agent error: ${err.message}`);
          }
        });

        setTimeout(() => {
          if (!resolved) {
            resolved = true;
            try { child.kill("SIGTERM"); } catch {}
            agent.status = "error";
            resolve(output.trim() + "\n(timed out after 3 minutes)");
          }
        }, timeout);
      } catch (err: any) {
        agent.status = "error";
        resolve(`Spawn error: ${err.message}`);
      }
    });
  }

  function killAgent(id: string): boolean {
    const agent = agents.get(id);
    if (!agent) return false;
    if (agent.process) {
      try { agent.process.kill("SIGTERM"); } catch {}
    }
    agent.status = "stopped";
    agent.process = null;
    agents.delete(id);
    return true;
  }

  function formatAgentList(): string {
    if (agents.size === 0) return "No active agents. Use /a2a spawn <role> to create one.";
    const lines: string[] = [];
    for (const [id, agent] of agents) {
      const uptime = Math.floor((Date.now() - agent.createdAt) / 1000);
      const uptimeStr = uptime > 60 ? `${Math.floor(uptime / 60)}m${uptime % 60}s` : `${uptime}s`;
      const statusIcon = agent.status === "running" ? "\u{1F7E2}" : agent.status === "idle" ? "\u{1F535}" : agent.status === "error" ? "\u{1F534}" : "\u26AB";
      lines.push(`  ${statusIcon} ${id} [${agent.role}] — ${agent.status} — uptime: ${uptimeStr}`);
    }
    return lines.join("\n");
  }

  pi.registerTool({
    name: "a2a_communicate",
    label: "Agent-to-Agent Communication",
    description: "Spawn specialized AI agents and communicate with them. Agents can be given roles like 'frontend', 'backend', 'tester', 'reviewer', 'devops', 'security', 'architect', or 'docs'.",
    parameters: Type.Object({
      action: Type.String({ description: "Action: 'spawn', 'send', 'broadcast', 'list', 'kill'" }),
      role: Type.Optional(Type.String({ description: "Agent role for spawn (e.g., 'frontend', 'backend', 'tester')" })),
      agent_id: Type.Optional(Type.String({ description: "Target agent ID for send/kill" })),
      message: Type.Optional(Type.String({ description: "Message to send to agent(s)" })),
    }),
    execute: async (args, ctx) => {
      const action = args.action.toLowerCase();

      switch (action) {
        case "spawn": {
          const role = args.role || "general";
          const agent = await spawnAgent(role, ctx.cwd);
          return { output: `\u2705 Agent spawned: ${agent.id} [${role}]\nSystem prompt: ${agent.systemPrompt.substring(0, 100)}...` };
        }
        case "send": {
          if (!args.agent_id) return { error: "Specify agent_id to send message to." };
          if (!args.message) return { error: "Specify message to send." };
          const agent = agents.get(args.agent_id);
          if (!agent) return { error: `Agent '${args.agent_id}' not found. Use action 'list' to see agents.` };
          if (agent.status === "running") return { error: `Agent '${args.agent_id}' is busy processing. Wait for completion.` };
          const response = await sendMessage(agent, args.message, ctx.cwd);
          return { output: `\u{1F4AC} Response from ${args.agent_id}:\n\n${response}` };
        }
        case "broadcast": {
          if (!args.message) return { error: "Specify message to broadcast." };
          if (agents.size === 0) return { error: "No active agents. Spawn some first." };
          const idle = Array.from(agents.values()).filter((a) => a.status === "idle");
          if (idle.length === 0) return { error: "All agents are busy. Wait for them to finish." };
          const promises = idle.map(async (agent) => {
            const response = await sendMessage(agent, args.message!, ctx.cwd);
            return `### ${agent.id} [${agent.role}]:\n${response.substring(0, 500)}`;
          });
          const results = await Promise.all(promises);
          return { output: `\u{1F4E2} Broadcast to ${idle.length} agents:\n\n${results.join("\n\n")}` };
        }
        case "list": {
          return { output: `\u{1F916} Active Agents:\n${formatAgentList()}` };
        }
        case "kill": {
          if (!args.agent_id) return { error: "Specify agent_id to kill." };
          const success = killAgent(args.agent_id);
          return { output: success ? `\u2705 Agent '${args.agent_id}' terminated.` : `\u274C Agent '${args.agent_id}' not found.` };
        }
        default:
          return { error: `Unknown action '${action}'. Use: spawn, send, broadcast, list, kill` };
      }
    },
  });

  pi.registerCommand("a2a", {
    description: "Agent-to-Agent: /a2a spawn <role> | list | send <id> <msg> | broadcast <msg> | kill <id>",
    handler: async (args, ctx) => {
      const parts = args.trim().split(/\s+/);
      const subcommand = (parts[0] || "list").toLowerCase();

      switch (subcommand) {
        case "spawn": {
          const role = parts[1] || "general";
          const validRoles = [...Object.keys(ROLE_PROMPTS), "general"];
          const agent = await spawnAgent(role, ctx.cwd);
          ctx.ui.notify(
            `\u{1F916} Agent Spawned!\n${'\u2500'.repeat(30)}\n\n` +
              `  ID:   ${agent.id}\n` +
              `  Role: ${role}\n` +
              `  Prompt: ${agent.systemPrompt.substring(0, 80)}...\n\n` +
              `Available roles: ${validRoles.join(", ")}\n` +
              `Send messages: /a2a send ${agent.id} <message>`,
            "info",
          );
          break;
        }
        case "list":
        case "ls": {
          ctx.ui.notify(`\u{1F916} Agent Registry\n${'\u2500'.repeat(30)}\n\n${formatAgentList()}`, "info");
          break;
        }
        case "send": {
          const agentId = parts[1];
          const message = parts.slice(2).join(" ");
          if (!agentId || !message) {
            ctx.ui.notify("Usage: /a2a send <agent-id> <message>", "warning");
            return;
          }
          const agent = agents.get(agentId);
          if (!agent) {
            ctx.ui.notify(`\u274C Agent '${agentId}' not found. Use /a2a list.`, "warning");
            return;
          }
          if (agent.status === "running") {
            ctx.ui.notify(`\u23F3 Agent '${agentId}' is busy. Wait for completion.`, "warning");
            return;
          }
          ctx.ui.notify(`\u{1F4E4} Sending to ${agentId}...`, "info");
          const response = await sendMessage(agent, message, ctx.cwd);
          ctx.ui.notify(`\u{1F4AC} ${agentId} responded:\n\n${response.substring(0, 2000)}`, "info");
          break;
        }
        case "broadcast": {
          const message = parts.slice(1).join(" ");
          if (!message) {
            ctx.ui.notify("Usage: /a2a broadcast <message>", "warning");
            return;
          }
          const idle = Array.from(agents.values()).filter((a) => a.status === "idle");
          if (idle.length === 0) {
            ctx.ui.notify("\u274C No idle agents. Spawn some first with /a2a spawn <role>.", "warning");
            return;
          }
          ctx.ui.notify(`\u{1F4E2} Broadcasting to ${idle.length} agents...`, "info");
          const promises = idle.map(async (agent) => {
            const response = await sendMessage(agent, message, ctx.cwd);
            return `\u{1F916} ${agent.id}:\n${response.substring(0, 500)}`;
          });
          const results = await Promise.all(promises);
          ctx.ui.notify(`\u{1F4E2} Broadcast Results:\n${'\u2500'.repeat(30)}\n\n${results.join("\n\n")}`, "info");
          break;
        }
        case "kill":
        case "stop": {
          const agentId = parts[1];
          if (!agentId) {
            ctx.ui.notify("Usage: /a2a kill <agent-id>", "warning");
            return;
          }
          if (agentId === "all") {
            const count = agents.size;
            for (const id of Array.from(agents.keys())) killAgent(id);
            ctx.ui.notify(`\u2705 Terminated ${count} agents.`, "info");
            return;
          }
          const success = killAgent(agentId);
          ctx.ui.notify(
            success ? `\u2705 Agent '${agentId}' terminated.` : `\u274C Agent '${agentId}' not found.`,
            success ? "info" : "warning",
          );
          break;
        }
        default:
          ctx.ui.notify(
            `\u{1F916} Agent-to-Agent Communication\n${'\u2500'.repeat(30)}\n\n` +
              "Commands:\n" +
              "  /a2a spawn <role>          - Spawn specialized agent\n" +
              "  /a2a list                  - List active agents\n" +
              "  /a2a send <id> <message>   - Send message to agent\n" +
              "  /a2a broadcast <message>   - Send to all idle agents\n" +
              "  /a2a kill <id|all>         - Terminate agent(s)\n\n" +
              `Roles: ${Object.keys(ROLE_PROMPTS).join(", ")}`,
            "info",
          );
      }
    },
  });
}