/**
 * remote-agent extension
 * ─────────────────────────────────────────────────────────────
 * Lets you talk to a friend's pi agent over Tailscale.
 *
 * Setup:
 *   Set REMOTE_AGENT_URL in your shell, e.g.:
 *     export REMOTE_AGENT_URL=https://friend-mac.tail1234.ts.net/rpc
 *   Or configure per-peer in ~/.pi/agent/remote-agents.json:
 *     { "friend": "https://friend-mac.tail1234.ts.net/rpc" }
 *
 * Usage inside pi:
 *   /ask-agent <prompt>              ← direct command
 *   /ask-agent @friend <prompt>      ← target a named peer
 *   The LLM can also call the ask_remote_agent tool directly.
 */

import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
import { Type } from "typebox";
import { readFileSync, existsSync } from "node:fs";
import { homedir } from "node:os";
import { join } from "node:path";

// ── Config ────────────────────────────────────────────────────────────────────

const PEERS_FILE = join(homedir(), ".pi", "agent", "remote-agents.json");

function loadPeers(): Record<string, string> {
  if (existsSync(PEERS_FILE)) {
    try {
      return JSON.parse(readFileSync(PEERS_FILE, "utf8"));
    } catch {
      return {};
    }
  }
  return {};
}

function resolveBaseUrl(peer?: string): string | null {
  if (peer) {
    const peers = loadPeers();
    return peers[peer] ?? null;
  }
  return process.env.REMOTE_AGENT_URL ?? null;
}

// ── Bridge Client ─────────────────────────────────────────────────────────────

interface AgentEvent {
  type: string;
  assistantMessageEvent?: { type: string; delta?: string };
  toolName?: string;
  args?: Record<string, unknown>;
  messages?: unknown[];
}

async function askRemoteAgent(
  baseUrl: string,
  prompt: string,
  onChunk: (text: string) => void,
  onTool: (name: string) => void,
): Promise<string> {
  // 1. Create session
  const createRes = await fetch(`${baseUrl}/sessions`, { method: "POST" });
  if (!createRes.ok) throw new Error(`Failed to create session: ${createRes.status}`);
  const { id } = (await createRes.json()) as { id: string };

  let fullText = "";
  let settled = false;

  // 2. Open SSE stream first
  const abortCtrl = new AbortController();

  const ssePromise = new Promise<void>(async (resolve, reject) => {
    try {
      const res = await fetch(`${baseUrl}/sessions/${id}/events`, {
        signal: abortCtrl.signal,
      });

      const reader = res.body!.getReader();
      const decoder = new TextDecoder();
      let buf = "";

      while (true) {
        const { done, value } = await reader.read();
        if (done) break;

        buf += decoder.decode(value, { stream: true });
        const lines = buf.split("\n");
        buf = lines.pop() ?? "";

        for (const line of lines) {
          if (!line.startsWith("data: ")) continue;
          const raw = line.slice(6).trim();
          if (!raw) continue;

          try {
            const ev = JSON.parse(raw) as AgentEvent;

            if (ev.type === "message_update") {
              const delta = ev.assistantMessageEvent;
              if (delta?.type === "text_delta" && delta.delta) {
                fullText += delta.delta;
                onChunk(delta.delta);
              }
            }

            if (ev.type === "tool_execution_start" && ev.toolName) {
              onTool(ev.toolName);
            }

            if (ev.type === "agent_end" || ev.type === "bridge_session_end") {
              if (!settled) {
                settled = true;
                abortCtrl.abort();
                resolve();
              }
              return;
            }
          } catch {
            // malformed line — skip
          }
        }
      }
      resolve();
    } catch (err: unknown) {
      // AbortError is expected when we cancel after agent_end
      if (err instanceof Error && err.name === "AbortError") {
        resolve();
      } else {
        reject(err);
      }
    }
  });

  // 3. Send prompt
  const promptRes = await fetch(`${baseUrl}/sessions/${id}/prompt`, {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ message: prompt }),
  });
  if (!promptRes.ok) throw new Error(`Failed to send prompt: ${promptRes.status}`);

  // 4. Wait for completion
  await ssePromise;

  // 5. Cleanup
  await fetch(`${baseUrl}/sessions/${id}`, { method: "DELETE" }).catch(() => {});

  return fullText;
}

// ── Extension ─────────────────────────────────────────────────────────────────

export default function (pi: ExtensionAPI) {

  // ── Tool: ask_remote_agent (LLM can call this) ──────────────────────────
  pi.registerTool({
    name: "ask_remote_agent",
    label: "Ask Remote Agent",
    description:
      "Send a task or question to a friend's pi agent over Tailscale and get the response. " +
      "Use this when you want to delegate work to another agent, get a second opinion, " +
      "or collaborate on a task. The remote agent has its own tools (read, bash, edit, write).",
    parameters: Type.Object({
      prompt: Type.String({
        description: "The task or question to send to the remote agent.",
      }),
      peer: Type.Optional(
        Type.String({
          description:
            "Named peer from ~/.pi/agent/remote-agents.json (e.g. 'friend'). " +
            "If omitted, uses REMOTE_AGENT_URL environment variable.",
        })
      ),
    }),

    async execute(_toolCallId, params, _signal, onUpdate, ctx) {
      const baseUrl = resolveBaseUrl(params.peer);

      if (!baseUrl) {
        return {
          content: [
            {
              type: "text",
              text:
                "❌ No remote agent URL configured.\n\n" +
                "Set REMOTE_AGENT_URL in your shell:\n" +
                "  export REMOTE_AGENT_URL=https://friend-mac.tail1234.ts.net/rpc\n\n" +
                "Or create ~/.pi/agent/remote-agents.json:\n" +
                '  { "friend": "https://friend-mac.tail1234.ts.net/rpc" }',
            },
          ],
          details: {},
        };
      }

      const peerLabel = params.peer ?? baseUrl;
      ctx.ui.setStatus("remote-agent", `🌙 Asking ${peerLabel}…`);

      let responseText = "";
      const toolsUsed: string[] = [];

      try {
        // Health check first
        const health = await fetch(`${baseUrl}/health`).catch(() => null);
        if (!health?.ok) {
          ctx.ui.setStatus("remote-agent", "");
          return {
            content: [
              {
                type: "text",
                text: `❌ Remote agent at ${baseUrl} is unreachable. Is the pi-bridge running?`,
              },
            ],
            details: {},
          };
        }

        responseText = await askRemoteAgent(
          baseUrl,
          params.prompt,
          (chunk) => {
            // Stream partial updates back to the TUI
            onUpdate({
              content: [{ type: "text", text: responseText + chunk }],
              details: { toolsUsed },
            });
          },
          (toolName) => {
            toolsUsed.push(toolName);
            ctx.ui.setStatus("remote-agent", `🔧 Remote: ${toolName}…`);
          }
        );

        ctx.ui.setStatus("remote-agent", "");

        return {
          content: [
            {
              type: "text",
              text: responseText || "(Remote agent returned no text response)",
            },
          ],
          details: { toolsUsed, peer: peerLabel },
        };
      } catch (err) {
        ctx.ui.setStatus("remote-agent", "");
        const msg = err instanceof Error ? err.message : String(err);
        return {
          content: [{ type: "text", text: `❌ Remote agent error: ${msg}` }],
          details: {},
        };
      }
    },
  });

  // ── Command: /ask-agent <prompt> or /ask-agent @peer <prompt> ────────────
  pi.registerCommand("ask-agent", {
    description: "Ask a remote pi agent via Tailscale. Usage: /ask-agent [@peer] <prompt>",
    handler: async (args, ctx) => {
      if (!args?.trim()) {
        ctx.ui.notify(
          "Usage: /ask-agent <prompt>  or  /ask-agent @friend <prompt>",
          "info"
        );
        return;
      }

      // Parse optional @peer prefix
      let peer: string | undefined;
      let prompt = args.trim();

      const peerMatch = prompt.match(/^@(\S+)\s+([\s\S]+)$/);
      if (peerMatch) {
        peer = peerMatch[1];
        prompt = peerMatch[2];
      }

      const baseUrl = resolveBaseUrl(peer);
      if (!baseUrl) {
        ctx.ui.notify(
          peer
            ? `Peer "@${peer}" not found in ~/.pi/agent/remote-agents.json`
            : "Set REMOTE_AGENT_URL or use /ask-agent @peer <prompt>",
          "warning"
        );
        return;
      }

      const peerLabel = peer ?? baseUrl;
      ctx.ui.notify(`Asking ${peerLabel}…`, "info");
      ctx.ui.setStatus("remote-agent", `🌙 Asking ${peerLabel}…`);

      try {
        const toolsUsed: string[] = [];
        const response = await askRemoteAgent(
          baseUrl,
          prompt,
          () => {}, // chunks handled by agent_end summary
          (toolName) => {
            toolsUsed.push(toolName);
            ctx.ui.setStatus("remote-agent", `🔧 Remote: ${toolName}…`);
          }
        );

        ctx.ui.setStatus("remote-agent", "");

        // Inject the result into the conversation as context
        await ctx.sendMessage(
          `## Response from remote agent (${peerLabel})\n\n` +
          `**Your prompt:** ${prompt}\n\n` +
          `**Response:**\n${response}` +
          (toolsUsed.length ? `\n\n*Tools used: ${toolsUsed.join(", ")}*` : "")
        );
      } catch (err) {
        ctx.ui.setStatus("remote-agent", "");
        const msg = err instanceof Error ? err.message : String(err);
        ctx.ui.notify(`Remote agent error: ${msg}`, "error");
      }
    },
  });
}