宜博 / lemonacp

/**
 * Chat 服务
 * 管理 WebSocket 会话、消息历史和多协议实例
 *
 * 支持的协议后端：
 * - claude: ClaudeStreamProtocol (stream-json)
 * - kimi:   AcpProtocol (kimi acp)
 * - gemini: AcpProtocol (gemini --experimental-acp)
 */

const { EventEmitter } = require('events');
const { randomUUID: uuidv4 } = require('crypto');
const fs = require('fs');
const path = require('path');
const { createConnection } = require('@mas/acp-core/transport/factory');

/**
 * 协议提供者配置
 */
const PROVIDERS = {
  claude: {
    label: 'Claude',
    command: 'claude',
    args: [
      '--input-format', 'stream-json',
      '--output-format', 'stream-json',
      '--verbose',
      '--include-partial-messages',
      '--dangerously-skip-permissions',
    ],
    factory: () => {
      const { ClaudeStreamProtocol } = require('@mas/acp-core');
      return new ClaudeStreamProtocol(
        { id: 'chat', command: 'claude', args: [] },
        { yolo: true, promptTimeout: 600000 }
      );
    },
  },
  kimi: {
    label: 'Kimi',
    command: 'kimi',
    args: ['acp'],
    factory: () => {
      const { AcpProtocol } = require('@mas/acp-core');
      return new AcpProtocol(
        { id: 'chat', command: 'kimi', args: ['acp'] },
        { yolo: true, promptTimeout: 0 }
      );
    },
  },
  gemini: {
    label: 'Gemini',
    command: 'gemini',
    args: ['--experimental-acp'],
    factory: () => {
      const { AcpProtocol } = require('@mas/acp-core');
      return new AcpProtocol(
        { id: 'chat', command: 'gemini', args: ['--experimental-acp'] },
        { yolo: true, promptTimeout: 0 }
      );
    },
  },
};

/**
 * Conversation — 持久化对话实体
 *
 * 一个 Conversation 对应一个 conversationId（URL 中使用），
 * 可跨多次 WebSocket 连接复用，实现多轮上下文对话。
 */
class ChatSession extends EventEmitter {
  constructor(sessionId, options = {}) {
    super();
    this.id = sessionId;       // conversationId（持久化标识）
    this.protocol = null;
    this.launcher = null;
    this.provider = null;
    this.messages = [];
    this.status = 'idle'; // idle, connecting, connected, error
    this.pid = null;
    this.createdAt = options.createdAt ? new Date(options.createdAt) : new Date();
    this.updatedAt = options.updatedAt ? new Date(options.updatedAt) : new Date();
    this.title = options.title || null;
    this.userId = options.userId || null; // 用户隔离
    this.options = options;

    // 持久化目录
    this._persistDir = options.persistDir || null;
  }

  /**
   * 是否已连接到 CLI 后端
   */
  get isConnected() {
    return this.protocol && this.status === 'connected';
  }

  /**
   * 获取历史消息摘要（用于注入 Claude 上下文）
   * 返回最近 N 轮对话的文本，供 CLI 续接上下文
   */
  getContextSummary(maxRounds = 10) {
    const recent = this.messages.slice(-maxRounds * 2);
    return recent.map(m => `${m.role}: ${(m.content || '').substring(0, 500)}`).join('\n');
  }

  /**
   * 连接到 CLI 后端
   * @param {string} workingDir 工作目录
   * @param {{ provider?: string }} options
   */
  async connect(workingDir, { provider = 'claude' } = {}) {
    if (this.protocol) {
      throw new Error('Session already connected');
    }

    const providerConfig = PROVIDERS[provider];
    if (!providerConfig) {
      throw new Error(`Unknown provider: ${provider}. Supported: ${Object.keys(PROVIDERS).join(', ')}`);
    }

    this.provider = provider;
    this.status = 'connecting';
    this.emit('statusChange', this.status);

    console.log(`[ChatSession:${this.id}] 使用 ${providerConfig.label} 后端连接...`);
    this.protocol = providerConfig.factory();

    // ── 统一事件通道：所有 ACP 事件走 acpEvent ──
    this.protocol.on('acp_event', (event) => {
      this.emit('acpEvent', event);
    });

    // ── 控制事件 ──
    this.protocol.on('error', (err) => {
      console.error(`[ChatSession:${this.id}] error: ${err.message}`);
      this.status = 'error';
      this.emit('statusChange', this.status);
      this.emit('error', err);
    });

    this.protocol.on('disconnected', () => {
      console.log(`[ChatSession:${this.id}] disconnected`);
      this.status = 'idle';
      this.pid = null;
      this.emit('statusChange', this.status);
      this.emit('disconnected');
    });

    try {
      // 通过 factory 创建 Launcher + Transport
      const agentDef = { transport: 'pipe' };
      const launcherDef = {
        type: 'process',
        command: providerConfig.command || 'claude',
        args: providerConfig.args || []
      };
      const { transport, launcher } = await createConnection(agentDef, launcherDef, workingDir);
      this.launcher = launcher;

      // 将 transport 交给 protocol
      await this.protocol.connect(transport, workingDir);
      // initialize 和 newSession 对 Claude 是 no-op，对 ACP 是必需的
      await this.protocol.initialize();
      await this.protocol.newSession({ cwd: workingDir });

      this.status = 'connected';
      this.pid = this.protocol.transport?.child?.pid
        || transport._process?.pid
        || null;
      this.updatedAt = new Date();
      console.log(`[ChatSession:${this.id}] connected via ${provider}, pid: ${this.pid}`);
      this.emit('statusChange', this.status);
      this.emit('connected', { pid: this.pid, provider });
    } catch (err) {
      console.error(`[ChatSession:${this.id}] connect failed:`, err.message);
      this.status = 'error';
      this.emit('statusChange', this.status);
      throw err;
    }
  }

  /**
   * 发送消息
   */
  async sendMessage(text) {
    if (!this.protocol || this.status !== 'connected') {
      throw new Error('Session not connected');
    }

    // 添加用户消息到历史
    this.addMessage({
      role: 'user',
      content: text,
      timestamp: new Date()
    });

    this.emit('thinking');

    try {
      const result = await this.protocol.sendPrompt(text);
      // 添加 AI 回复到历史
      this.addMessage({
        role: 'assistant',
        content: result.result || '',
        is_error: result.is_error || false,
        duration_ms: result.duration_ms,
        timestamp: new Date()
      });
      this.updatedAt = new Date();
      this.emit('messageComplete', result);
      return result;
    } catch (err) {
      this.emit('error', err);
      throw err;
    }
  }

  /**
   * 添加消息到历史
   */
  addMessage(message) {
    const msg = {
      id: uuidv4(),
      timestamp: new Date(),
      ...message
    };
    this.messages.push(msg);
    this.updatedAt = new Date();

    // 自动设置 title（第一条用户消息）
    if (!this.title && msg.role === 'user' && msg.content) {
      this.title = msg.content.substring(0, 80);
    }

    // 持久化到磁盘
    this._appendMessage(msg);
    this._saveMeta();

    this.emit('messageAdded', msg);
    return msg;
  }

  /**
   * 断开连接
   */
  disconnect() {
    if (this.protocol) {
      this.protocol.disconnect();
      this.protocol = null;
    }
    if (this.launcher) {
      try { this.launcher.stop(); } catch { }
      this.launcher = null;
    }
    this.status = 'idle';
    this.pid = null;
    this.provider = null;
    this._saveMeta();
    this.emit('statusChange', this.status);
  }

  /**
   * 获取会话信息
   */
  getInfo() {
    return {
      id: this.id,
      status: this.status,
      provider: this.provider,
      pid: this.pid,
      messageCount: this.messages.length,
      createdAt: this.createdAt,
      updatedAt: this.updatedAt,
      title: this.title,
      userId: this.userId,
    };
  }

  /**
   * 获取消息历史
   */
  getMessages() {
    return [...this.messages];
  }

  // ═══════════════════════════════════════════════════
  //  持久化
  // ═══════════════════════════════════════════════════

  /** 确保持久化目录存在 */
  _ensurePersistDir() {
    if (!this._persistDir) return false;
    if (!fs.existsSync(this._persistDir)) {
      fs.mkdirSync(this._persistDir, { recursive: true });
    }
    return true;
  }

  /** 保存会话元数据 */
  _saveMeta() {
    if (!this._ensurePersistDir()) return;
    const meta = {
      id: this.id,
      title: this.title,
      provider: this.provider,
      status: this.status,
      messageCount: this.messages.length,
      createdAt: this.createdAt,
      updatedAt: this.updatedAt,
      userId: this.userId,
    };
    try {
      fs.writeFileSync(
        path.join(this._persistDir, 'meta.json'),
        JSON.stringify(meta, null, 2),
        'utf-8'
      );
    } catch (err) {
      console.warn(`[ChatSession:${this.id}] 保存 meta 失败:`, err.message);
    }
  }

  /** 追加消息到 messages.jsonl */
  _appendMessage(msg) {
    if (!this._ensurePersistDir()) return;
    try {
      fs.appendFileSync(
        path.join(this._persistDir, 'messages.jsonl'),
        JSON.stringify(msg) + '\n',
        'utf-8'
      );
    } catch (err) {
      console.warn(`[ChatSession:${this.id}] 追加消息失败:`, err.message);
    }
  }
}

/**
 * Chat 服务管理器
 */
class ClaudeChatService extends EventEmitter {
  constructor(conversationsDir = null) {
    super();
    this.sessions = new Map();
    this.conversationsDir = conversationsDir
      || path.resolve(__dirname, '../../conversations');
  }

  /**
   * 获取支持的 provider 列表
   */
  getProviders() {
    return Object.entries(PROVIDERS).map(([id, config]) => ({
      id,
      label: config.label,
    }));
  }

  /**
   * 创建新对话
   * @param {string|null} convId  conversationId
   * @param {{ userId?: string|number }} options
   */
  createSession(convId = null, options = {}) {
    const id = convId || uuidv4();

    if (this.sessions.has(id)) {
      throw new Error(`Conversation ${id} already exists`);
    }

    // 按 userId 隔离存储路径
    const userId = options.userId ? String(options.userId) : null;
    const baseDir = userId
      ? path.join(this.conversationsDir, `user-${userId}`)
      : this.conversationsDir;
    const persistDir = path.join(baseDir, `conv-${id}`);
    const session = new ChatSession(id, { ...options, userId, persistDir });
    this.sessions.set(id, session);

    // 创建目录 + 保存初始 meta
    session._ensurePersistDir();
    session._saveMeta();

    // 转发会话事件
    session.on('connected', (data) => this.emit('sessionConnected', { sessionId: id, ...data }));
    session.on('disconnected', () => this.emit('sessionDisconnected', { sessionId: id }));
    session.on('error', (err) => this.emit('sessionError', { sessionId: id, error: err }));
    session.on('messageAdded', (msg) => this.emit('messageAdded', { sessionId: id, message: msg }));

    this.emit('sessionCreated', { sessionId: id });

    return session;
  }

  /**
   * 获取或创建对话（核心方法 — 支持 WS 断线重连）
   * 如果 conversationId 已存在 → 返回已有 Conversation
   * 如果不存在 → 创建新 Conversation
   */
  getOrCreateConversation(convId, options = {}) {
    const existing = this.sessions.get(convId);
    if (existing) {
      // 校验用户归属
      const userId = options.userId ? String(options.userId) : null;
      if (existing.userId && userId && existing.userId !== userId) {
        throw new Error('Conversation belongs to another user');
      }
      console.log(`[Service] 复用已有对话 ${convId} (messages: ${existing.messages.length}, status: ${existing.status})`);
      return existing;
    }
    // 不存在 → 创建新的
    console.log(`[Service] 创建新对话 ${convId}`);
    return this.createSession(convId, options);
  }

  /**
   * 获取会话
   */
  getSession(sessionId) {
    return this.sessions.get(sessionId);
  }

  /**
   * 删除会话
   */
  deleteSession(sessionId) {
    const session = this.sessions.get(sessionId);
    if (session) {
      session.disconnect();
      this.sessions.delete(sessionId);
      this.emit('sessionDeleted', { sessionId });
      return true;
    }
    return false;
  }

  /**
   * 获取所有会话列表
   */
  getSessions() {
    return Array.from(this.sessions.values()).map(s => s.getInfo());
  }

  /**
   * 获取指定用户的会话列表
   */
  getSessionsByUser(userId) {
    const uid = String(userId);
    return Array.from(this.sessions.values())
      .filter(s => s.userId === uid)
      .map(s => s.getInfo());
  }

  /**
   * 清理已断开的会话
   */
  cleanup() {
    for (const [id, session] of this.sessions) {
      if (session.status === 'idle' || session.status === 'error') {
        this.deleteSession(id);
      }
    }
  }

  // ═══════════════════════════════════════════════════
  //  磁盘恢复
  // ═══════════════════════════════════════════════════

  /**
   * 从 conversations/ 目录恢复历史会话
   * 扫描所有 chatbot-* 子目录
   */
  loadFromDisk() {
    if (!fs.existsSync(this.conversationsDir)) return;

    let loaded = 0;

    // 辅助函数: 从目录恢复对话
    const restoreFromDir = (baseDir, userId) => {
      if (!fs.existsSync(baseDir)) return;
      // 兼容旧目录 chatbot- 和新目录 conv-
      const convDirs = fs.readdirSync(baseDir, { withFileTypes: true })
        .filter(d => d.isDirectory() && (d.name.startsWith('conv-') || d.name.startsWith('chatbot-')));

      for (const dir of convDirs) {
        const convId = dir.name.replace(/^(conv-|chatbot-)/, '');
        if (this.sessions.has(convId)) continue;

        const dirPath = path.join(baseDir, dir.name);
        const metaPath = path.join(dirPath, 'meta.json');
        const msgsPath = path.join(dirPath, 'messages.jsonl');

        if (!fs.existsSync(metaPath)) continue;

        try {
          const meta = JSON.parse(fs.readFileSync(metaPath, 'utf-8'));
          const session = new ChatSession(convId, {
            persistDir: dirPath,
            createdAt: meta.createdAt,
            updatedAt: meta.updatedAt,
            title: meta.title,
            userId: meta.userId || userId || null,
          });
          session.provider = meta.provider || null;

          // 恢复消息历史
          if (fs.existsSync(msgsPath)) {
            const lines = fs.readFileSync(msgsPath, 'utf-8')
              .split('\n')
              .filter(l => l.trim());
            for (const line of lines) {
              try {
                session.messages.push(JSON.parse(line));
              } catch { /* skip malformed line */ }
            }
          }

          this.sessions.set(convId, session);
          loaded++;
        } catch (err) {
          console.warn(`[ClaudeChatService] 恢复对话 ${convId} 失败:`, err.message);
        }
      }
    };

    // 1. 扫描 user-* 子目录下的 chatbot 会话
    const topEntries = fs.readdirSync(this.conversationsDir, { withFileTypes: true });
    for (const entry of topEntries) {
      if (entry.isDirectory() && entry.name.startsWith('user-')) {
        const userId = entry.name.replace('user-', '');
        restoreFromDir(path.join(this.conversationsDir, entry.name), userId);
      }
    }

    // 2. 兼容: 扫描根目录下的旧 chatbot 会话
    restoreFromDir(this.conversationsDir, null);

    if (loaded > 0) {
      console.log(`[ClaudeChatService] 从磁盘恢复了 ${loaded} 个历史会话`);
    }
  }
}

// 单例实例 — 启动时自动从磁盘恢复历史会话
const service = new ClaudeChatService();
service.loadFromDisk();

module.exports = {
  ClaudeChatService,
  ChatSession,
  service,
  PROVIDERS
};