宜博 / lemonboss

/**
 * ═══════════════════════════════════════════════════════════════
 *  AgentRegistry — Agent 注册中心
 * ═══════════════════════════════════════════════════════════════
 *
 *  职责:
 *    1. Agent 的 CRUD (注册/查看/修改/删除)
 *    2. 生命周期管理 (创建目录 → 加载 → 启动 → 停止 → 销毁)
 *    3. 组织架构树维护 (parent-child 关系)
 *    4. 与 AgentBus 协调 (创建时注册到总线, 删除时注销)
 *
 *  用法:
 *    const registry = new AgentRegistry({
 *      agentsDir: 'data/agents',
 *      db: sqliteDb,
 *      bus: agentBus,
 *    });
 *
 *    // 创建
 *    const agent = await registry.create({ type: 'cxo', name: 'CEO' });
 *
 *    // 查看
 *    const all = registry.list();
 *    const ceo = registry.get('cxo-ceo-001');
 *
 *    // 修改
 *    await registry.updateBootstrap('cxo-ceo-001', 'agent.md', newContent);
 *
 *    // 删除
 *    await registry.destroy('cxo-ceo-001');
 */

import { readFile, writeFile, rm, access, constants } from 'node:fs/promises';
import { join } from 'node:path';
import { randomUUID } from 'node:crypto';
import { EventEmitter } from 'node:events';
import { Agent } from './Agent.js';
import { AgentScaffold } from './AgentScaffold.js';

export class AgentRegistry extends EventEmitter {

  /**
   * @param {Object} options
   * @param {string} options.agentsDir - Agent 实例根目录 (如 data/agents)
   * @param {Object} [options.db]      - SQLite 数据库实例 (better-sqlite3)
   * @param {Object} [options.bus]     - AgentBus 消息总线实例
   */
  constructor(options) {
    super();
    this.agentsDir = options.agentsDir;
    this.db = options.db || null;
    this.bus = options.bus || null;
    this.scaffold = new AgentScaffold(this.agentsDir);

    // 内存中的 Agent 实例缓存
    this.agents = new Map(); // agentId → Agent 实例

    if (this.db) {
      this._initDB();
    }
  }

  // ─────────────────────────────────────────
  //  CREATE — 创建 Agent
  // ─────────────────────────────────────────

  /**
   * 创建一个新的 Agent
   *
   * @param {Object} options
   * @param {string} options.type      - 'master' | 'cxo' | 'role' | 'skill'
   * @param {string} options.name      - 显示名称
   * @param {string} [options.id]      - 自定义 ID (默认自动生成)
   * @param {string} [options.parentId]  - 上级 Agent ID
   * @param {string} [options.goalId]    - 绑定的 Goal ID
   * @param {Object} [options.bootstrap] - 自定义 Bootstrap Files 内容
   *
   * @returns {Agent} 创建并加载的 Agent 实例
   */
  async create(options) {
    const id = options.id || `${options.type}-${this._shortUUID()}`;

    // 1. 用 Scaffold 生成目录
    const agentDir = await this.scaffold.create({
      id,
      type: options.type,
      name: options.name,
      parentId: options.parentId,
      goalId: options.goalId,
      bootstrap: options.bootstrap,
    });

    // 2. 注册到数据库
    if (this.db) {
      this.db.prepare(`
        INSERT INTO agents (id, name, type, parent_id, goal_id, status, dir_path)
        VALUES (?, ?, ?, ?, ?, 'idle', ?)
      `).run(id, options.name, options.type, options.parentId || null, options.goalId || null, agentDir);
    }

    // 3. 加载 Agent 实例
    const agent = new Agent(agentDir);
    await agent.load();
    this.agents.set(id, agent);

    // 4. 注册到消息总线
    if (this.bus) {
      this.bus.register(id, {
        type: options.type,
        parentId: options.parentId,
      }, (msg) => this._handleMessage(id, msg));
    }

    this.emit('agent:created', agent.getSummary());
    return agent;
  }

  // ─────────────────────────────────────────
  //  READ — 查看 Agent
  // ─────────────────────────────────────────

  /**
   * 获取单个 Agent
   * @returns {Agent|null}
   */
  get(agentId) {
    return this.agents.get(agentId) || null;
  }

  /**
   * 列出所有 Agent
   * @param {Object} [filter] - 过滤条件
   * @param {string} [filter.type]     - 按类型过滤
   * @param {string} [filter.parentId] - 按上级过滤
   * @param {string} [filter.status]   - 按状态过滤
   *
   * @returns {Object[]} Agent 摘要列表
   */
  list(filter = {}) {
    // 优先从数据库查
    if (this.db) {
      let sql = 'SELECT * FROM agents WHERE 1=1';
      const params = [];

      if (filter.type) {
        sql += ' AND type = ?';
        params.push(filter.type);
      }
      if (filter.parentId) {
        sql += ' AND parent_id = ?';
        params.push(filter.parentId);
      }
      if (filter.status) {
        sql += ' AND status = ?';
        params.push(filter.status);
      }

      sql += ' ORDER BY created_at';
      return this.db.prepare(sql).all(...params);
    }

    // 无数据库时从内存读
    let results = Array.from(this.agents.values()).map(a => a.getSummary());

    if (filter.type) results = results.filter(a => a.type === filter.type);
    if (filter.parentId) results = results.filter(a => a.parentId === filter.parentId);

    return results;
  }

  /**
   * 获取组织架构树
   * @returns {Object} 树状结构 { agent, children: [...] }
   */
  getOrgTree() {
    const allAgents = this.list();
    const byId = new Map(allAgents.map(a => [a.id, { ...a, children: [] }]));

    let root = null;
    for (const agent of byId.values()) {
      if (agent.parent_id && byId.has(agent.parent_id)) {
        byId.get(agent.parent_id).children.push(agent);
      } else if (!agent.parent_id) {
        root = agent;
      }
    }

    return root || { children: allAgents };
  }

  // ─────────────────────────────────────────
  //  UPDATE — 修改 Agent
  // ─────────────────────────────────────────

  /**
   * 更新 Agent 的 Bootstrap File
   * (这是"管理"的核心操作 — 上层 Agent 修改下层的 agent.md / soul.md)
   *
   * @param {string} agentId
   * @param {string} filename - Bootstrap 文件名 (如 'agent.md', 'soul.md')
   * @param {string} content  - 新的文件内容
   */
  async updateBootstrap(agentId, filename, content) {
    const agent = this.agents.get(agentId);
    if (!agent) throw new Error(`Agent 不存在: ${agentId}`);

    // 写入文件
    const filepath = join(agent.dir, filename);
    await writeFile(filepath, content, 'utf-8');

    // 重新加载 Agent (热更新)
    await agent.load();

    // 更新数据库时间戳
    if (this.db) {
      this.db.prepare('UPDATE agents SET updated_at = CURRENT_TIMESTAMP WHERE id = ?').run(agentId);
    }

    this.emit('agent:updated', { agentId, filename });
    return agent;
  }

  /**
   * 更新 Agent 运行配置
   */
  async updateConfig(agentId, configPatch) {
    const agent = this.agents.get(agentId);
    if (!agent) throw new Error(`Agent 不存在: ${agentId}`);

    // 读取现有 config
    const configPath = join(agent.dir, 'config.json');
    const existing = JSON.parse(await readFile(configPath, 'utf-8'));

    // 合并
    const merged = { ...existing, ...configPatch };
    await writeFile(configPath, JSON.stringify(merged, null, 2), 'utf-8');

    // 重新加载
    await agent.load();

    this.emit('agent:config-updated', { agentId });
    return agent;
  }

  /**
   * 更新 Agent 状态
   */
  async updateStatus(agentId, status) {
    if (this.db) {
      this.db.prepare('UPDATE agents SET status = ?, updated_at = CURRENT_TIMESTAMP WHERE id = ?')
        .run(status, agentId);
    }

    const agent = this.agents.get(agentId);
    if (agent) agent.status = status;

    this.emit('agent:status', { agentId, status });
  }

  // ─────────────────────────────────────────
  //  DELETE — 删除 Agent
  // ─────────────────────────────────────────

  /**
   * 销毁一个 Agent (停止 + 从总线注销 + 删除目录 + 从数据库移除)
   *
   * @param {string} agentId
   * @param {Object} [options]
   * @param {boolean} [options.recursive=false] - 是否同时销毁所有下级 Agent
   */
  async destroy(agentId, options = {}) {
    const { recursive = false } = options;

    // 递归删除下级
    if (recursive) {
      const children = this.list({ parentId: agentId });
      for (const child of children) {
        await this.destroy(child.id, { recursive: true });
      }
    }

    const agent = this.agents.get(agentId);

    // 1. 停止 Agent
    if (agent && agent.status === 'running') {
      await agent.stop();
    }

    // 2. 从消息总线注销
    if (this.bus) {
      this.bus.unregister(agentId);
    }

    // 3. 删除目录
    if (agent) {
      await rm(agent.dir, { recursive: true, force: true });
    }

    // 4. 从数据库移除
    if (this.db) {
      this.db.prepare('DELETE FROM agents WHERE id = ?').run(agentId);
    }

    // 5. 从内存缓存移除
    this.agents.delete(agentId);

    this.emit('agent:destroyed', agentId);
  }

  // ─────────────────────────────────────────
  //  生命周期: 启动 / 停止
  // ─────────────────────────────────────────

  /**
   * 启动一个 Agent (加载 + 启动 Loop)
   */
  async start(agentId) {
    const agent = this.agents.get(agentId);
    if (!agent) throw new Error(`Agent 不存在: ${agentId}`);

    await agent.start();
    await this.updateStatus(agentId, 'running');

    return agent;
  }

  /**
   * 停止一个 Agent
   */
  async stop(agentId) {
    const agent = this.agents.get(agentId);
    if (!agent) throw new Error(`Agent 不存在: ${agentId}`);

    await agent.stop();
    await this.updateStatus(agentId, 'idle');

    return agent;
  }

  /**
   * 初始化: 扫描 data/agents/ 目录, 加载所有已存在的 Agent
   */
  async loadAll() {
    const { readdir } = await import('node:fs/promises');
    const entries = await readdir(this.agentsDir, { withFileTypes: true });
    const loaded = [];

    for (const entry of entries) {
      if (!entry.isDirectory()) continue;

      const agentDir = join(this.agentsDir, entry.name);
      const configPath = join(agentDir, 'config.json');

      try {
        await access(configPath, constants.R_OK);
      } catch {
        continue; // 不是合法的 Agent 目录
      }

      const agent = new Agent(agentDir);
      await agent.load();
      this.agents.set(agent.id, agent);

      // 注册到消息总线
      if (this.bus) {
        this.bus.register(agent.id, {
          type: agent.type,
          parentId: agent.parentId,
        }, (msg) => this._handleMessage(agent.id, msg));
      }

      loaded.push(agent.id);
    }

    this.emit('registry:loaded', loaded);
    return loaded;
  }

  // ─────────────────────────────────────────
  //  内部方法
  // ─────────────────────────────────────────

  /**
   * 处理收到的消息 (转发给 Agent 实例)
   */
  _handleMessage(agentId, msg) {
    const agent = this.agents.get(agentId);
    if (agent) {
      agent.emit('message', msg);
    }
  }

  /**
   * 初始化数据库表
   */
  _initDB() {
    this.db.exec(`
      CREATE TABLE IF NOT EXISTS agents (
        id          TEXT PRIMARY KEY,
        name        TEXT NOT NULL,
        type        TEXT NOT NULL CHECK(type IN ('master', 'cxo', 'role', 'skill')),
        parent_id   TEXT,
        goal_id     TEXT,
        status      TEXT DEFAULT 'idle'
                    CHECK(status IN ('idle', 'running', 'paused', 'error')),
        dir_path    TEXT,
        created_at  DATETIME DEFAULT CURRENT_TIMESTAMP,
        updated_at  DATETIME DEFAULT CURRENT_TIMESTAMP
      );

      CREATE INDEX IF NOT EXISTS idx_agents_type ON agents(type);
      CREATE INDEX IF NOT EXISTS idx_agents_parent ON agents(parent_id);
      CREATE INDEX IF NOT EXISTS idx_agents_status ON agents(status);
    `);
  }

  /**
   * 生成短 UUID
   */
  _shortUUID() {
    return randomUUID().split('-')[0];
  }
}

export default AgentRegistry;