#!/usr/bin/env node
// federation-agent/server/registry.js · v1.0
// 服务器端 · Agent 注册管理中心
// 挂载到 gatekeeper 上，管理所有联邦 Agent 节点
// ICE-GL-ZY001 · D129 · 2026-06-10

const WebSocket = require("ws");
const http = require("http");

// ===== 配置 =====
const PORT = process.env.AGENT_PORT || 3912;
const AUTH_TOKEN = process.env.AGENT_AUTH || "zy_gtw_federation_agent_v1";

// ===== 注册表 =====
class AgentRegistry {
  constructor() {
    this.agents = new Map();      // agentId → { ws, info, status }
    this.personas = new Map();     // personaId → [agentId, ...]
    this.users = new Map();        // userId → [agentId, ...]
    this.commandQueue = new Map(); // agentId → [pending commands]
  }

  register(ws, info) {
    const { agentId, personaId, userId, hostname, platform, workDir } = info;

    const agent = {
      ws,
      info: { agentId, personaId, userId, hostname, platform, workDir, version: info.version || "1.0.0" },
      status: "online",
      connectedAt: new Date().toISOString(),
      lastHeartbeat: Date.now(),
      commandCount: 0,
    };

    this.agents.set(agentId, agent);

    if (personaId) {
      if (!this.personas.has(personaId)) this.personas.set(personaId, []);
      this.personas.get(personaId).push(agentId);
    }
    if (userId) {
      if (!this.users.has(userId)) this.users.set(userId, []);
      this.users.get(userId).push(agentId);
    }

    console.log(`[REGISTRY] + ${agentId} (${personaId || "未注册"} @ ${userId || "未知用户"} · ${hostname})`);
    return agent;
  }

  unregister(agentId) {
    const agent = this.agents.get(agentId);
    if (!agent) return;

    this.agents.delete(agentId);

    if (agent.info.personaId) {
      const list = this.personas.get(agent.info.personaId);
      if (list) {
        const idx = list.indexOf(agentId);
        if (idx > -1) list.splice(idx, 1);
      }
    }
    if (agent.info.userId) {
      const list = this.users.get(agent.info.userId);
      if (list) {
        const idx = list.indexOf(agentId);
        if (idx > -1) list.splice(idx, 1);
      }
    }

    console.log(`[REGISTRY] - ${agentId} (断开)`);
  }

  sendCommand(agentId, command) {
    const agent = this.agents.get(agentId);
    if (!agent || !agent.ws || agent.ws.readyState !== WebSocket.OPEN) {
      // 加入待发送队列
      if (!this.commandQueue.has(agentId)) this.commandQueue.set(agentId, []);
      this.commandQueue.get(agentId).push(command);
      return { ok: false, error: "agent offline" };
    }

    agent.ws.send(JSON.stringify({ type: "command", command }));
    agent.commandCount++;
    return { ok: true };
  }

  flushQueue(agentId) {
    const queue = this.commandQueue.get(agentId);
    if (!queue || queue.length === 0) return 0;

    let sent = 0;
    while (queue.length > 0) {
      const cmd = queue.shift();
      const result = this.sendCommand(agentId, cmd);
      if (result.ok) sent++;
      else break;
    }
    return sent;
  }

  broadcast(personaIds, command) {
    const results = [];
    for (const personaId of personaIds) {
      const agentIds = this.personas.get(personaId) || [];
      for (const agentId of agentIds) {
        results.push({ agentId, ...this.sendCommand(agentId, command) });
      }
    }
    return results;
  }

  list(target) {
    if (target === "all") {
      return Array.from(this.agents.values()).map(a => ({
        agentId: a.info.agentId,
        personaId: a.info.personaId,
        userId: a.info.userId,
        hostname: a.info.hostname,
        status: a.status,
        commandCount: a.commandCount,
        connectedAt: a.connectedAt,
      }));
    }
    // 按用户过滤
    const agentIds = this.users.get(target) || [];
    return agentIds.map(id => {
      const a = this.agents.get(id);
      return a ? {
        agentId: a.info.agentId,
        personaId: a.info.personaId,
        status: a.status,
        commandCount: a.commandCount,
      } : null;
    }).filter(Boolean);
  }

  healthCheck() {
    const now = Date.now();
    for (const [agentId, agent] of this.agents) {
      if (now - agent.lastHeartbeat > 120000) {
        agent.status = "timeout";
        console.log(`[REGISTRY] ⚠ ${agentId} 心跳超时`);
      }
    }
  }
}

// ===== WebSocket 服务器 =====
const registry = new AgentRegistry();
const server = http.createServer((req, res) => {
  // REST API 端点
  if (req.url === "/api/agents" && req.method === "GET") {
    const target = new URL(req.url, "http://localhost").searchParams.get("user") || "all";
    res.writeHead(200, { "Content-Type": "application/json" });
    res.end(JSON.stringify(registry.list(target)));
    return;
  }
  if (req.url === "/api/command" && req.method === "POST") {
    let body = "";
    req.on("data", chunk => body += chunk);
    req.on("end", () => {
      try {
        const { agentId, command } = JSON.parse(body);
        const result = registry.sendCommand(agentId, command);
        res.writeHead(result.ok ? 200 : 404, { "Content-Type": "application/json" });
        res.end(JSON.stringify(result));
      } catch (e) {
        res.writeHead(400, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ ok: false, error: e.message }));
      }
    });
    return;
  }
  res.writeHead(404);
  res.end("Not Found");
});

const wss = new WebSocket.Server({ server, path: "/agent" });

wss.on("connection", (ws, req) => {
  const url = new URL(req.url, "http://localhost");
  const agentId = url.searchParams.get("agentId") || "unknown";
  const personaId = url.searchParams.get("persona") || "";
  const userId = url.searchParams.get("user") || "";

  console.log(`[WS] 新连接: ${agentId}`);

  // 先不发 welcome，等 register 消息
  ws.agentId = agentId;

  ws.on("message", (data) => {
    try {
      const msg = JSON.parse(data.toString());

      switch (msg.type) {
        case "register": {
          const agent = registry.register(ws, msg);
          ws.send(JSON.stringify({
            type: "welcome",
            message: `光湖联邦托管 · 注册成功 · ${agent.info.personaId || agentId}`,
            personaId: agent.info.personaId,
            timestamp: Date.now(),
          }));
          // 发送队列中的命令
          const flushed = registry.flushQueue(agentId);
          if (flushed > 0) console.log(`[REGISTRY] → ${agentId} 补发 ${flushed} 条待处理指令`);
          break;
        }
        case "result": {
          console.log(`[RESULT] ${agentId}: ${msg.command} → ${msg.ok ? "OK" : "FAIL"}`);
          if (msg.stderr) console.log(`  stderr: ${msg.stderr.substring(0, 200)}`);
          break;
        }
        case "heartbeat": {
          const agent = registry.agents.get(agentId);
          if (agent) agent.lastHeartbeat = Date.now();
          break;
        }
        case "health_report": {
          console.log(`[HEALTH] ${agentId}: uptime=${(msg.uptime / 60).toFixed(1)}min · mem=${(msg.memory.heapUsed / 1024 / 1024).toFixed(1)}MB`);
          break;
        }
        case "confirm_request": {
          console.log(`[CONFIRM] ${agentId}: 请求确认危险操作: ${msg.operation}`);
          break;
        }
        default:
          console.log(`[MSG] ${agentId}: ${msg.type}`);
      }
    } catch (e) {
      console.error(`[ERROR] ${agentId}: ${e.message}`);
    }
  });

  ws.on("close", () => {
    registry.unregister(agentId);
  });

  ws.on("error", (err) => {
    console.error(`[WS ERROR] ${agentId}: ${err.message}`);
  });
});

// 健康检查定时器
setInterval(() => registry.healthCheck(), 30000);

server.listen(PORT, () => {
  console.log(`[REGISTRY] 光湖联邦 Agent 注册中心 · 端口 ${PORT}`);
  console.log(`[REGISTRY] WebSocket: ws://0.0.0.0:${PORT}/agent`);
  console.log(`[REGISTRY] REST API:  http://0.0.0.0:${PORT}/api/agents`);
});

process.on("SIGINT", () => {
  console.log("[REGISTRY] 关闭中...");
  wss.close();
  server.close();
  process.exit(0);
});