宜博 / lemonboss

/**
 * Chat WebSocket 路由处理器
 *
 * 统一 WS 消息格式：
 * - connected:     连接确认
 * - initialized:   初始化完成 (含 provider)
 * - thinking:      开始思考
 * - acp_event:     所有内容事件（文本/工具/权限等）
 * - result:        轮次结束
 * - error:         错误
 * - disconnected:  断开
 */

const { service } = require('../services/claude-chat-service');
const { extractUserFromToken, COOKIE_NAME } = require('../middleware/jwt-auth');
const config = require('../config');
const fs = require('fs');
const path = require('path');

/**
 * WebSocket 消息处理器
 */
class ClaudeChatWSHandler {
  constructor(ws, wsId, user) {
    this.ws = ws;
    this.wsId = wsId;         // 临时 WS 连接 ID
    this.session = null;      // 关联的 Conversation (ChatSession)
    this.user = user;
    this.activeSkills = new Map();
    this._skillsChanged = false;
    this._injectedSkillsHash = '';

    this.ws.on('message', (data) => this.handleMessage(data));
    this.ws.on('close', () => this.handleClose());
    this.ws.on('error', (err) => this.handleError(err));
  }

  /**
   * 发送消息给客户端
   */
  send(type, payload = {}) {
    if (this.ws.readyState === 1) { // OPEN
      this.ws.send(JSON.stringify({ type, ...payload }));
    }
  }

  /**
   * 处理收到的消息
   */
  async handleMessage(data) {
    try {
      const message = JSON.parse(data.toString());
      const { type, ...payload } = message;

      switch (type) {
        case 'init':
          await this.handleInit(payload);
          break;

        case 'message':
          await this.handleSendMessage(payload);
          break;

        case 'activate_skill':
          this.handleActivateSkill(payload);
          break;

        case 'deactivate_skill':
          this.handleDeactivateSkill(payload);
          break;

        case 'disconnect':
          await this.handleDisconnect();
          break;

        case 'getHistory':
          this.handleGetHistory();
          break;

        case 'getProviders':
          this.handleGetProviders();
          break;

        default:
          this.send('error', { error: 'Unknown message type' });
      }
    } catch (err) {
      console.error('WebSocket message error:', err);
      this.send('error', { error: err.message });
    }
  }

  /**
   * 初始化会话 — 接收 conversationId 复用已有对话
   */
  async handleInit({ workingDir = process.cwd(), provider = 'claude', conversationId } = {}) {
    try {
      // 如果客户端没传 conversationId，生成一个
      const convId = conversationId || `conv-${Date.now().toString(36)}-${Math.random().toString(36).substr(2, 6)}`;

      // 获取或创建 Conversation
      this.session = service.getOrCreateConversation(convId, {
        userId: this.user.id
      });

      // 重新绑定事件到当前 WS
      this.bindSessionEvents();

      // 如果 Conversation 已有连接的 CLI，复用它
      if (this.session.isConnected) {
        console.log(`[WS:${this.wsId}] 复用已有 CLI 连接 (conv: ${convId}, pid: ${this.session.pid})`);
      } else {
        // 需要新建 CLI 连接
        await this.session.connect(workingDir, { provider });

        // 如果有历史消息，注入上下文摘要让 Claude 记住之前的对话
        if (this.session.messages.length > 0) {
          const contextSummary = this.session.getContextSummary();
          if (contextSummary) {
            console.log(`[WS:${this.wsId}] 注入历史上下文 (${this.session.messages.length} msgs)`);
            const contextMsg = `以下是之前的对话历史摘要，请你记住这些上下文，然后等待用户的下一条消息：\n\n${contextSummary}\n\n---\n请回复"已了解，请继续"即可。`;
            await this.session.sendMessage(contextMsg);
          }
        }
      }

      this.send('initialized', {
        conversationId: convId,
        sessionId: convId,   // 兼容旧前端
        pid: this.session.pid,
        provider: this.session.provider,
        messageCount: this.session.messages.length
      });
    } catch (err) {
      this.send('error', { error: err.message });
    }
  }

  /**
   * 绑定 Conversation 事件到当前 WS
   * 每次新的 WS 连接都重新绑定，确保事件发到正确的 WS
   */
  bindSessionEvents() {
    // 移除旧的事件监听器（避免重复绑定）
    this.session.removeAllListeners('acpEvent');
    this.session.removeAllListeners('thinking');
    this.session.removeAllListeners('messageComplete');
    // 保留 error/disconnected/connected 的内部监听

    // ── 唯一的内容通道 ──
    this.session.on('acpEvent', (event) => {
      this.send('acp_event', { event });
    });

    // ── 控制事件 ──
    this.session.on('thinking', () => {
      console.log(`[WS:${this.wsId}] thinking...`);
      this.send('thinking');
    });

    this.session.on('messageComplete', (result) => {
      console.log(`[WS:${this.wsId}] result:`, JSON.stringify(result)?.substring(0, 120));
      this.send('result', { result });
    });
  }

  /**
   * 激活 Skill — 加载 SKILL.md + 关键参考文件
   */
  handleActivateSkill({ skillName }) {
    if (!skillName) {
      this.send('error', { error: 'skillName is required' });
      return;
    }

    const skillDir = path.join(config.skillsDataDir, skillName);
    const skillMdPath = path.join(skillDir, 'SKILL.md');
    if (!fs.existsSync(skillMdPath)) {
      this.send('error', { error: `Skill "${skillName}" not found` });
      return;
    }

    // 加载完整 skill bundle（SKILL.md + 参考文件）
    const bundle = this._loadSkillBundle(skillName, skillDir);
    this.activeSkills.set(skillName, bundle);
    this._skillsChanged = true;

    console.log(`[WS:${this.wsId}] activated skill: ${skillName} (${bundle.length} chars)`);
    this.send('skill_activated', {
      skillName,
      activeSkills: [...this.activeSkills.keys()],
    });
  }

  /**
   * 加载 Skill Bundle：SKILL.md + digest_summary.md + knowhow.md
   */
  _loadSkillBundle(skillName, skillDir) {
    const parts = [];

    // 1. 主文件 SKILL.md（必须）
    const skillMd = fs.readFileSync(path.join(skillDir, 'SKILL.md'), 'utf-8');
    parts.push(skillMd);

    // 2. 参考文件（可选，按优先级加载）
    const refFiles = [
      { path: 'digest/digest_summary.md', label: 'Digest Summary' },
      { path: 'hunter/knowhow.md', label: 'Domain Knowhow' },
      { path: 'digest/benchmark.md', label: 'Quality Benchmark' },
    ];

    for (const ref of refFiles) {
      const refPath = path.join(skillDir, ref.path);
      if (fs.existsSync(refPath)) {
        try {
          const content = fs.readFileSync(refPath, 'utf-8');
          // 限制单个参考文件不超过 8000 字符，防止 token 爆炸
          const truncated = content.length > 8000
            ? content.substring(0, 8000) + '\n\n... (truncated)'
            : content;
          parts.push(`\n### ${ref.label}\n${truncated}`);
        } catch { /* skip unreadable files */ }
      }
    }

    return parts.join('\n\n');
  }

  /**
   * 取消激活 Skill
   */
  handleDeactivateSkill({ skillName }) {
    this.activeSkills.delete(skillName);
    this._skillsChanged = true;
    console.log(`[WS:${this.wsId}] deactivated skill: ${skillName}`);
    this.send('skill_deactivated', {
      skillName,
      activeSkills: [...this.activeSkills.keys()],
    });
  }

  /**
   * 构建 Skill 前缀（仅在 skills 变更时注入）
   * @returns {string} 需要注入的前缀，空字符串表示不需要注入
   */
  _buildSkillPrefix() {
    // 无 skill → 不注入
    if (this.activeSkills.size === 0) return '';

    // skills 没变化 → 不重复注入（AI 有上下文记忆）
    const currentHash = [...this.activeSkills.keys()].sort().join(',');
    if (!this._skillsChanged && this._injectedSkillsHash === currentHash) {
      return '';
    }

    // 标记为已注入
    this._injectedSkillsHash = currentHash;
    this._skillsChanged = false;

    const parts = [];
    for (const [name, content] of this.activeSkills) {
      parts.push(`<skill name="${name}">\n${content}\n</skill>`);
    }

    return '请严格按照以下 Skill 指导来完成任务：\n\n'
      + parts.join('\n\n')
      + '\n\n---\n\n';
  }

  /**
   * 构建附件上下文（将上传文件内容注入到消息中）
   */
  _buildAttachmentContext(attachments) {
    if (!attachments || attachments.length === 0) return '';

    const parts = [];
    for (const att of attachments) {
      const filePath = att.storedPath;
      if (!filePath || !fs.existsSync(filePath)) continue;

      try {
        if (att.fileType && att.fileType.startsWith('image/')) {
          // 图片 → base64 引用
          const b64 = fs.readFileSync(filePath, 'base64');
          const ext = path.extname(att.fileName).slice(1) || 'png';
          parts.push(
            `<attachment name="${att.fileName}" type="image">\n` +
            `[Image: ${att.fileName} (${(att.fileSize / 1024).toFixed(1)}KB, ${ext})]\n` +
            `Base64: data:${att.fileType};base64,${b64}\n` +
            `</attachment>`
          );
        } else {
          // 文本/代码/CSV → 直接读取内容
          let content = fs.readFileSync(filePath, 'utf-8');
          // 限制单文件 20000 字符
          if (content.length > 20000) {
            content = content.substring(0, 20000) + '\n\n... (truncated, total ' + content.length + ' chars)';
          }
          parts.push(
            `<attachment name="${att.fileName}" type="${att.fileType}">\n` +
            `${content}\n` +
            `</attachment>`
          );
        }
      } catch (e) {
        console.warn(`[WS:${this.wsId}] Failed to read attachment: ${att.fileName}`, e.message);
      }
    }

    if (parts.length === 0) return '';
    return '\n\n以下是用户上传的附件内容：\n\n' + parts.join('\n\n') + '\n\n---\n\n';
  }

  /**
   * 发送消息（智能 Skill + 附件注入）
   */
  async handleSendMessage({ text, skills, attachments }) {
    if (!this.session) {
      this.send('error', { error: 'Session not initialized' });
      return;
    }

    // 如果消息附带了 skills 列表，先激活它们
    if (skills && Array.isArray(skills)) {
      for (const sk of skills) {
        if (!this.activeSkills.has(sk)) {
          this.handleActivateSkill({ skillName: sk });
        }
      }
    }

    // 构建完整消息: [skill 前缀] + 用户文本 + [附件上下文]
    const skillPrefix = this._buildSkillPrefix();
    const attachContext = this._buildAttachmentContext(attachments);
    const fullText = (skillPrefix || '') + text + (attachContext || '');

    try {
      await this.session.sendMessage(fullText);
    } catch (err) {
      this.send('error', { error: err.message });
    }
  }

  /**
   * 断开连接
   */
  async handleDisconnect() {
    if (this.session) {
      this.session.disconnect();
    }
    this.send('disconnected');
  }

  /**
   * 获取消息历史
   */
  handleGetHistory() {
    if (!this.session) {
      this.send('error', { error: 'Session not initialized' });
      return;
    }

    this.send('history', {
      messages: this.session.getMessages()
    });
  }

  /**
   * 获取支持的 provider 列表
   */
  handleGetProviders() {
    this.send('providers', {
      providers: service.getProviders()
    });
  }

  /**
   * 处理连接关闭 — 不销毁 Conversation，允许重连
   */
  handleClose() {
    console.log(`WebSocket ${this.wsId} closed (conv: ${this.session?.id || 'none'})`);
    // 不断开 Conversation，允许通过新 WS 重新连接
  }

  /**
   * 处理错误
   */
  handleError(err) {
    console.error(`WebSocket error for ${this.wsId}:`, err);
  }
}

/**
 * 初始化 WebSocket 服务器
 */
function initClaudeChatWSServer(wss) {
  wss.on('connection', (ws, req) => {
    // 从 WS handshake Cookie 中提取 JWT 鉴权
    const cookieHeader = req.headers.cookie || '';
    const cookies = {};
    cookieHeader.split(';').forEach(c => {
      const [k, ...v] = c.trim().split('=');
      if (k) cookies[k] = v.join('=');
    });
    const token = cookies[COOKIE_NAME];
    const user = extractUserFromToken(token);

    if (!user) {
      console.warn('[WS] Unauthorized WebSocket connection rejected');
      ws.close(4001, 'Unauthorized');
      return;
    }

    // wsId 是临时的 WS 连接标识，不是 conversationId
    const wsId = 'ws-' + Date.now().toString(36) + Math.random().toString(36).substr(2, 5);
    console.log(`WebSocket connected: ${wsId} (user: ${user.id})`);

    const handler = new ClaudeChatWSHandler(ws, wsId, user);

    // 发送连接成功消息（客户端收到后发送 init + conversationId）
    handler.send('connected', { wsId });
  });

  console.log('Chat WebSocket server initialized');
}

module.exports = {
  ClaudeChatWSHandler,
  initClaudeChatWSServer
};