宜博 / lemonboss

/**
 * PromptBuilder - 根据 Agent 类型拼接 protocol + prompt
 */
const fs = require('fs');
const path = require('path');

class PromptBuilder {
  constructor(config, baseDir) {
    this.config = config;
    this.baseDir = baseDir;
    this.protocolCache = new Map();
    this.promptCache = new Map();
  }

  /**
   * 动态生成环境快照 (Company, Agent, Goals)
   */
  getSystemStateXML() {
    let state = '<env_status>\n';
    state += `  <timestamp>${new Date().toISOString()}</timestamp>\n`;
    try {
      const { CompanyService } = require('../../services/company-service');
      const { GoalsService } = require('../../services/goals');
      
      const cs = new CompanyService();
      const companies = cs.list();
      
      if (companies.length > 0) {
        state += '  <companies>\n';
        for (const c of companies) {
          state += `    <company id="${c.id}" name="${c.name}">\n`;
          const agents = cs.listAgents(c.id);
          if (agents && agents.length > 0) {
             for (const a of agents) {
               state += `      <agent id="${c.id}:${a.id}" role="${a.type}" status="${a.status || 'ready'}" />\n`;
             }
          }
          const gs = new GoalsService(c.id);
          const goals = gs.list();
          if (goals && goals.length > 0) {
             state += '      <active_goals>\n';
             for (const g of goals.filter(g => g.status !== 'done')) {
                state += `        <goal id="${g.id}" title="${g.title}" status="${g.status}" />\n`;
             }
             state += '      </active_goals>\n';
          }
          state += `    </company>\n`;
        }
        state += '  </companies>\n';
      }

      const globalGs = new GoalsService();
      const globalGoals = globalGs.list();
      if (globalGoals && globalGoals.length > 0) {
        state += '  <global_active_goals>\n';
        for (const g of globalGoals.filter(g => g.status !== 'done')) {
           state += `    <goal id="${g.id}" title="${g.title}" status="${g.status}" />\n`;
        }
        state += '  </global_active_goals>\n';
      }
    } catch (e) {
      state += `  <!-- Error loading env state: ${e.message} -->\n`;
    }
    state += '</env_status>\n';
    return state;
  }

  /**
   * 加载文件内容（带缓存）
   */
  loadFile(filePath) {
    const fullPath = path.isAbsolute(filePath)
      ? filePath
      : path.join(this.baseDir, filePath);

    if (!this.promptCache.has(fullPath)) {
      if (!fs.existsSync(fullPath)) {
        throw new Error(`File not found: ${fullPath}`);
      }
      this.promptCache.set(fullPath, fs.readFileSync(fullPath, 'utf-8'));
    }

    return this.promptCache.get(fullPath);
  }

  /**
   * 获取 role protocol 内容
   */
  getProtocol(type) {
    const protocolPath = this.config.rolePrompts?.[type] || this.config.protocols?.[type];
    if (!protocolPath) {
      return '';
    }
    return this.loadFile(protocolPath);
  }

  /**
   * 获取 agent prompt 内容
   */
  getAgentPrompt(agentName) {
    const agent = this.config.agents[agentName];
    if (!agent) {
      throw new Error(`Agent not found: ${agentName}`);
    }

    const promptPath = agent.prompt || agent.systemPrompt;
    if (!promptPath) {
      return '';
    }

    return this.loadFile(promptPath);
  }

  /**
   * 构建可用 Agent 列表（用于 master protocol）
   */
  buildAvailableAgentsList() {
    const excluded = new Set(['master', 'gateway']);
    return Object.entries(this.config.agents)
      .filter(([name]) => !excluded.has(name))
      .map(([name, cfg]) => `- **${name}** (${cfg.type}): ${cfg.description}`)
      .join('\n');
  }

  /**
   * 构建完整的 system prompt
   * - SKILL.md 模式: 直接使用 SKILL.md（已包含 protocol），不再叠加 protocol
   * - 旧模式: protocol + agent prompt（向后兼容非 Innovator agent）
   */
  build(agentName, userTask = '') {
    const agent = this.config.agents[agentName];
    if (!agent) {
      throw new Error(`Agent not found: ${agentName}`);
    }

    const parts = [];
    const promptPath = agent.prompt || agent.systemPrompt || '';
    const isSkillBased = promptPath.endsWith('SKILL.md');

    if (isSkillBased) {
      // SKILL.md 已包含 protocol，直接使用
      const skillPrompt = this.getAgentPrompt(agentName);
      if (skillPrompt) parts.push(skillPrompt);
    } else {
      // 旧模式：先加 protocol，再加 agent prompt
      const protocol = this.getProtocol(agent.type);
      if (protocol) {
        let processedProtocol = protocol;
        if (agent.type === 'master') {
          processedProtocol = processedProtocol.replace(
            '{{AVAILABLE_AGENTS}}',
            this.buildAvailableAgentsList()
          );
        }
        parts.push(processedProtocol);
      }
      const agentPrompt = this.getAgentPrompt(agentName);
      if (agentPrompt) parts.push(agentPrompt);
      if (agent.companyContextPath && fs.existsSync(agent.companyContextPath)) {
        parts.push(this.loadFile(agent.companyContextPath));
      }
    }

    // User Task
    if (userTask) {
      parts.push(`---\n\n## 当前任务\n\n${userTask}`);
    }

    // 在最终拼接前，注入上帝视角的环境感知
    parts.unshift(this.getSystemStateXML());

    let result = parts.join('\n\n---\n\n');

    // 替换 {{AVAILABLE_AGENTS}} 占位符（可能在 SKILL.md 或 protocol 中）
    if (result.includes('{{AVAILABLE_AGENTS}}')) {
      result = result.replace(/\{\{AVAILABLE_AGENTS\}\}/g, this.buildAvailableAgentsList());
    }

    return result;
  }
}

module.exports = { PromptBuilder };