ai-terminal-assistant/packages/core/src/checkpoint/manager.ts

/**
 * 检查点管理器
 * 作为编排器，委托具体工作给各个子模块
 */

import * as path from 'path';
import { getCheckpointsDir } from '../constants/paths.js';
import { ShadowGit, createShadowGit } from './shadow-git.js';
import { CheckpointLock } from './lock.js';
import { CheckpointSafetyChecker } from './safety.js';
import { WorkspacePathValidator } from './path-validator.js';
import { CommitMessageGenerator } from './commit-message.js';
import { LFSPatternLoader } from './lfs.js';
import {
  type CheckpointMetadata,
  type CheckpointConfig,
  type CheckpointTrigger,
  type RollbackOptions,
  type RollbackResult,
  type DiffInfo,
  type FileDiff,
  type CheckpointEventListener,
  type RollbackRecord,
  type UnrevertResult,
  type SafetyCheckResult,
} from './types.js';

// 子模块
import { CheckpointStore, type CreateCheckpointOptions } from './checkpoint-store.js';
import { CheckpointRollback } from './checkpoint-rollback.js';
import { CheckpointSession } from './checkpoint-session.js';
import { CheckpointEvents } from './checkpoint-events.js';

/**
 * 检查点管理器
 */
export class CheckpointManager {
  private shadowGit: ShadowGit;
  private config: CheckpointConfig;
  private workDir: string;
  private initialized = false;
  private lastCheckpointTime = 0;

  // 子模块
  private store: CheckpointStore;
  private rollbackHandler: CheckpointRollback;
  private session: CheckpointSession;
  private events: CheckpointEvents;

  // 辅助组件
  private lock: CheckpointLock;
  private safetyChecker: CheckpointSafetyChecker;
  private pathValidator: WorkspacePathValidator;
  private lfsLoader: LFSPatternLoader;

  // 防止重复创建检查点的最小间隔 (毫秒)
  private static readonly MIN_CHECKPOINT_INTERVAL = 1000;

  constructor(workDir: string, config: Partial<CheckpointConfig> = {}) {
    this.workDir = path.resolve(workDir);
    this.config = {
      enabled: true,
      autoCheckpoint: {
        beforeWrite: true,
        beforeEdit: true,
        beforeDelete: true,
        beforeMove: true,
        beforeBash: false,
      },
      maxCheckpoints: 100,
      maxAge: 7 * 24 * 60 * 60 * 1000,
      storageDir: getCheckpointsDir(),
      ...config,
    };

    this.shadowGit = createShadowGit(this.workDir, this.config.storageDir);

    // 初始化辅助组件
    this.lock = new CheckpointLock(this.shadowGit.getShadowGitDir());
    this.safetyChecker = new CheckpointSafetyChecker(this.workDir);
    this.pathValidator = new WorkspacePathValidator();
    this.lfsLoader = new LFSPatternLoader();

    // 初始化子模块
    const commitMessageGenerator = new CommitMessageGenerator();
    this.store = new CheckpointStore(this.shadowGit, this.lock, commitMessageGenerator, this.config);
    this.rollbackHandler = new CheckpointRollback(this.shadowGit, this.lock, this.safetyChecker, this.store);
    this.session = new CheckpointSession(this.store);
    this.events = new CheckpointEvents();
  }

  /**
   * 初始化检查点管理器
   */
  async initialize(): Promise<void> {
    if (this.initialized) return;

    if (!this.config.enabled) {
      this.initialized = true;
      return;
    }

    // 验证工作目录路径
    const pathValidation = this.pathValidator.validate(this.workDir);
    if (!pathValidation.valid) {
      throw new Error(`Invalid workspace path: ${pathValidation.reason}`);
    }

    // 加载 LFS 模式
    await this.lfsLoader.loadPatterns(this.workDir);

    // 初始化 Shadow Git
    await this.shadowGit.initialize();

    // 加载检查点索引
    await this.store.loadIndex();

    this.initialized = true;
  }

  // ============================================================================
  // 检查点操作（委托给 store）
  // ============================================================================

  /**
   * 判断是否应该为指定工具创建检查点
   */
  shouldCreateCheckpoint(tool: string): boolean {
    return this.store.shouldCreateForTool(tool);
  }

  /**
   * 在工具执行前创建检查点
   */
  async beforeToolExecution(
    tool: string,
    params: Record<string, unknown>
  ): Promise<string | null> {
    if (!this.shouldCreateCheckpoint(tool)) {
      return null;
    }

    // 防止过于频繁的检查点创建
    const now = Date.now();
    if (now - this.lastCheckpointTime < CheckpointManager.MIN_CHECKPOINT_INTERVAL) {
      return null;
    }

    try {
      const checkpoint = await this.createCheckpoint({
        trigger: `tool:${tool}` as CheckpointTrigger,
        toolCall: { tool, params },
        description: this.store.generateDescription(tool, params),
      });

      this.lastCheckpointTime = now;
      return checkpoint.id;
    } catch (error) {
      console.warn('Failed to create checkpoint:', error);
      return null;
    }
  }

  /**
   * 创建检查点
   */
  async createCheckpoint(options: CreateCheckpointOptions): Promise<CheckpointMetadata> {
    await this.initialize();

    if (!this.config.enabled) {
      throw new Error('Checkpoint system is disabled');
    }

    const checkpoint = await this.store.create(
      options,
      this.session.getCurrentSessionId(),
      (id) => this.session.recordCheckpoint(id)
    );

    // 触发事件
    this.events.emitCreated(checkpoint);

    // 异步清理
    this.cleanupAsync();

    return checkpoint;
  }

  /**
   * 创建命名检查点
   */
  async createNamedCheckpoint(name: string, description?: string): Promise<CheckpointMetadata> {
    return this.createCheckpoint({
      name,
      description,
      trigger: 'manual',
    });
  }

  /**
   * 获取所有检查点
   */
  async listCheckpoints(): Promise<CheckpointMetadata[]> {
    await this.initialize();
    return this.store.list();
  }

  /**
   * 获取指定检查点
   */
  async getCheckpoint(idOrHash: string): Promise<CheckpointMetadata | null> {
    await this.initialize();
    return this.store.get(idOrHash);
  }

  /**
   * 获取最近的检查点
   */
  async getLatestCheckpoint(): Promise<CheckpointMetadata | null> {
    await this.initialize();
    return this.store.getLatest();
  }

  /**
   * 删除检查点
   */
  async deleteCheckpoint(checkpointId: string): Promise<boolean> {
    await this.initialize();

    const checkpoint = this.store.delete(checkpointId);
    if (checkpoint) {
      this.events.emitDeleted(checkpoint);
      return true;
    }
    return false;
  }

  // ============================================================================
  // 差异操作
  // ============================================================================

  /**
   * 获取检查点与当前工作区的差异
   */
  async getDiff(checkpointId: string): Promise<DiffInfo> {
    await this.initialize();

    const checkpoint = await this.getCheckpoint(checkpointId);
    if (!checkpoint) {
      throw new Error(`Checkpoint not found: ${checkpointId}`);
    }

    return this.shadowGit.getDiffSummary(checkpoint.commitHash, 'HEAD');
  }

  /**
   * 获取两个检查点之间的差异
   */
  async getDiffBetween(fromId: string, toId: string): Promise<DiffInfo> {
    await this.initialize();

    const fromCheckpoint = await this.getCheckpoint(fromId);
    const toCheckpoint = await this.getCheckpoint(toId);

    if (!fromCheckpoint) {
      throw new Error(`Checkpoint not found: ${fromId}`);
    }
    if (!toCheckpoint) {
      throw new Error(`Checkpoint not found: ${toId}`);
    }

    return this.shadowGit.getDiffSummary(
      fromCheckpoint.commitHash,
      toCheckpoint.commitHash
    );
  }

  /**
   * 获取文件差异详情
   */
  async getFileDiff(checkpointId: string, filePath: string): Promise<FileDiff> {
    await this.initialize();

    const checkpoint = await this.getCheckpoint(checkpointId);
    if (!checkpoint) {
      throw new Error(`Checkpoint not found: ${checkpointId}`);
    }

    const head = await this.shadowGit.getHead();
    return this.shadowGit.getFileDiff(checkpoint.commitHash, head, filePath);
  }

  // ============================================================================
  // 回滚操作（委托给 rollbackHandler）
  // ============================================================================

  /**
   * 回滚到检查点
   */
  async rollback(options: RollbackOptions): Promise<RollbackResult> {
    await this.initialize();

    return this.rollbackHandler.rollback(options, (event) => {
      if (event.type === 'restored' && event.checkpoint) {
        this.events.emitRestored(event.checkpoint, event.details as {
          files: string[];
          previousCommit: string;
          mode: string;
        });
      }
    });
  }

  /**
   * 撤销操作（回滚到上一个检查点）
   */
  async undo(): Promise<RollbackResult> {
    await this.initialize();

    return this.rollbackHandler.undo((event) => {
      if (event.type === 'restored' && event.checkpoint) {
        this.events.emitRestored(event.checkpoint, event.details as {
          files: string[];
          previousCommit: string;
          mode: string;
        });
      }
    });
  }

  /**
   * 撤销最近一次回滚
   */
  async unrevert(): Promise<UnrevertResult> {
    await this.initialize();
    return this.rollbackHandler.unrevert();
  }

  /**
   * 检查是否可以 unrevert
   */
  canUnrevert(): boolean {
    return this.rollbackHandler.canUnrevert();
  }

  /**
   * 获取最后一次回滚记录
   */
  getLastRollback(): RollbackRecord | null {
    return this.rollbackHandler.getLastRollback();
  }

  /**
   * 执行安全检查
   */
  async checkSafety(checkpointId: string): Promise<SafetyCheckResult> {
    await this.initialize();

    const checkpoint = await this.getCheckpoint(checkpointId);
    if (!checkpoint) {
      return {
        safe: false,
        warnings: [],
        errors: ['Checkpoint not found'],
      };
    }

    return this.safetyChecker.checkBeforeRollback(checkpoint, {
      getCheckpoint: (id: string) => Promise.resolve(this.store.get(id)),
      listCheckpoints: () => Promise.resolve(this.store.list()),
      getDiff: (id: string) => this.getDiff(id),
    });
  }

  // ============================================================================
  // 会话操作（委托给 session）
  // ============================================================================

  /**
   * 开始新会话
   */
  async startSession(sessionId?: string): Promise<string> {
    await this.initialize();

    return this.session.startSession(sessionId, (event) => {
      if (event.checkpoint) {
        this.events.emitCreated(event.checkpoint);
      }
    });
  }

  /**
   * 结束当前会话
   */
  async endSession(): Promise<void> {
    return this.session.endSession((event) => {
      if (event.checkpoint) {
        this.events.emitCreated(event.checkpoint);
      }
    });
  }

  /**
   * 获取当前会话 ID
   */
  getCurrentSessionId(): string | null {
    return this.session.getCurrentSessionId();
  }

  /**
   * 获取会话的所有检查点
   */
  async getSessionCheckpoints(sessionId: string): Promise<CheckpointMetadata[]> {
    await this.initialize();
    return this.session.getSessionCheckpoints(sessionId);
  }

  /**
   * 创建与消息关联的检查点
   */
  async createMessageCheckpoint(
    messageId: string,
    turnIndex?: number,
    options?: {
      trigger?: CheckpointTrigger;
      description?: string;
    }
  ): Promise<CheckpointMetadata> {
    await this.initialize();

    return this.session.createMessageCheckpoint(messageId, turnIndex, options, (event) => {
      this.events.emitCreated(event.checkpoint);
    });
  }

  /**
   * 获取与消息关联的检查点
   */
  async getMessageCheckpoints(messageId: string): Promise<CheckpointMetadata[]> {
    await this.initialize();
    return this.session.getMessageCheckpoints(messageId);
  }

  /**
   * 撤销整个会话的修改
   */
  async undoSession(sessionId: string): Promise<RollbackResult> {
    const sessionCheckpoints = await this.getSessionCheckpoints(sessionId);
    if (sessionCheckpoints.length === 0) {
      throw new Error(`No checkpoints found for session: ${sessionId}`);
    }

    const startCheckpoint = this.session.getSessionStartCheckpoint(sessionId);

    if (!startCheckpoint) {
      return this.rollback({ target: sessionCheckpoints[0].id });
    }

    return this.rollback({ target: startCheckpoint.id });
  }

  // ============================================================================
  // 事件操作（委托给 events）
  // ============================================================================

  /**
   * 添加事件监听器
   */
  addEventListener(listener: CheckpointEventListener): void {
    this.events.addEventListener(listener);
  }

  /**
   * 移除事件监听器
   */
  removeEventListener(listener: CheckpointEventListener): void {
    this.events.removeEventListener(listener);
  }

  // ============================================================================
  // 清理操作
  // ============================================================================

  /**
   * 异步清理过期检查点
   */
  private async cleanupAsync(): Promise<void> {
    setTimeout(async () => {
      try {
        await this.cleanup();
      } catch (error) {
        console.warn('Checkpoint cleanup failed:', error);
      }
    }, 100);
  }

  /**
   * 清理过期检查点
   */
  async cleanup(): Promise<number> {
    await this.initialize();

    const checkpoints = this.store.list();
    const now = Date.now();
    let deletedCount = 0;

    // 按时间过期清理
    for (const checkpoint of checkpoints) {
      if (now - checkpoint.timestamp > this.config.maxAge) {
        if (this.store.delete(checkpoint.id)) {
          this.events.emitDeleted(checkpoint);
          deletedCount++;
        }
      }
    }

    // 按数量限制清理
    const remaining = checkpoints.length - deletedCount;
    if (remaining > this.config.maxCheckpoints) {
      const toDelete = checkpoints.slice(this.config.maxCheckpoints);
      for (const checkpoint of toDelete) {
        if (this.store.delete(checkpoint.id)) {
          this.events.emitDeleted(checkpoint);
          deletedCount++;
        }
      }
    }

    if (deletedCount > 0) {
      this.events.emitCleanup(deletedCount);
      await this.shadowGit.cleanup(this.config.maxCheckpoints);
    }

    return deletedCount;
  }

  // ============================================================================
  // 其他方法
  // ============================================================================

  /**
   * 获取检查点存储统计
   */
  async getStats(): Promise<{
    count: number;
    oldestTimestamp: number | null;
    newestTimestamp: number | null;
  }> {
    const checkpoints = await this.listCheckpoints();

    return {
      count: checkpoints.length,
      oldestTimestamp: checkpoints.length > 0
        ? checkpoints[checkpoints.length - 1].timestamp
        : null,
      newestTimestamp: checkpoints.length > 0 ? checkpoints[0].timestamp : null,
    };
  }

  /**
   * 检查是否启用
   */
  isEnabled(): boolean {
    return this.config.enabled;
  }

  /**
   * 获取配置
   */
  getConfig(): CheckpointConfig {
    return { ...this.config };
  }

  /**
   * 获取 LFS 模式加载器
   */
  getLfsLoader(): LFSPatternLoader {
    return this.lfsLoader;
  }

  /**
   * 检查文件是否由 LFS 管理
   */
  isLfsFile(filePath: string): boolean {
    return this.lfsLoader.isLfsFile(filePath);
  }
}

// 全局检查点管理器实例
let globalCheckpointManager: CheckpointManager | null = null;

/**
 * 获取全局检查点管理器实例
 */
export function getCheckpointManager(): CheckpointManager {
  if (!globalCheckpointManager) {
    globalCheckpointManager = new CheckpointManager(process.cwd());
  }
  return globalCheckpointManager;
}

/**
 * 初始化全局检查点管理器
 */
export async function initCheckpointManager(
  workDir: string,
  config?: Partial<CheckpointConfig>
): Promise<CheckpointManager> {
  globalCheckpointManager = new CheckpointManager(workDir, config);
  await globalCheckpointManager.initialize();
  return globalCheckpointManager;
}

/**
 * 重置全局检查点管理器 (用于测试)
 */
export function resetCheckpointManager(): void {
  globalCheckpointManager = null;
}