tsdoc/ts/context/enhanced-context.ts

import * as plugins from '../plugins.js';
import type { ContextMode, IContextResult, IFileInfo, TaskType, IFileMetadata } from './types.js';
import { ContextTrimmer } from './context-trimmer.js';
import { ConfigManager } from './config-manager.js';
import { LazyFileLoader } from './lazy-file-loader.js';
import { ContextCache } from './context-cache.js';
import { ContextAnalyzer } from './context-analyzer.js';

/**
 * Enhanced ProjectContext that supports context optimization strategies
 */
export class EnhancedContext {
  private projectDir: string;
  private trimmer: ContextTrimmer;
  private configManager: ConfigManager;
  private lazyLoader: LazyFileLoader;
  private cache: ContextCache;
  private analyzer: ContextAnalyzer;
  private contextMode: ContextMode = 'trimmed';
  private tokenBudget: number = 190000; // Default for o4-mini
  private contextResult: IContextResult = {
    context: '',
    tokenCount: 0,
    includedFiles: [],
    trimmedFiles: [],
    excludedFiles: [],
    tokenSavings: 0
  };
  
  /**
   * Create a new EnhancedContext
   * @param projectDirArg The project directory
   */
  constructor(projectDirArg: string) {
    this.projectDir = projectDirArg;
    this.configManager = ConfigManager.getInstance();
    this.trimmer = new ContextTrimmer(this.configManager.getTrimConfig());
    this.lazyLoader = new LazyFileLoader(projectDirArg);
    this.cache = new ContextCache(projectDirArg, this.configManager.getCacheConfig());
    this.analyzer = new ContextAnalyzer(
      projectDirArg,
      this.configManager.getPrioritizationWeights(),
      this.configManager.getTierConfig()
    );
  }
  
  /**
   * Initialize the context builder
   */
  public async initialize(): Promise<void> {
    await this.configManager.initialize(this.projectDir);
    this.tokenBudget = this.configManager.getMaxTokens();
    this.trimmer.updateConfig(this.configManager.getTrimConfig());
    await this.cache.init();
  }
  
  /**
   * Set the context mode
   * @param mode The context mode to use
   */
  public setContextMode(mode: ContextMode): void {
    this.contextMode = mode;
  }
  
  /**
   * Set the token budget
   * @param maxTokens The maximum tokens to use
   */
  public setTokenBudget(maxTokens: number): void {
    this.tokenBudget = maxTokens;
  }
  
  /**
   * Gather files from the project
   * @param includePaths Optional paths to include
   * @param excludePaths Optional paths to exclude
   */
  public async gatherFiles(includePaths?: string[], excludePaths?: string[]): Promise<Record<string, plugins.smartfile.SmartFile | plugins.smartfile.SmartFile[]>> {
    const smartfilePackageJSON = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'package.json'),
      this.projectDir,
    );
    
    const smartfilesReadme = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'readme.md'),
      this.projectDir,
    );

    const smartfilesReadmeHints = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'readme.hints.md'),
      this.projectDir,
    );
    
    const smartfilesNpmextraJSON = await plugins.smartfile.SmartFile.fromFilePath(
      plugins.path.join(this.projectDir, 'npmextra.json'),
      this.projectDir,
    );
    
    // Use provided include paths or default to all TypeScript files
    const includeGlobs = includePaths?.map(path => `${path}/**/*.ts`) || ['ts*/**/*.ts'];
    
    // Get TypeScript files
    const smartfilesModPromises = includeGlobs.map(glob => 
      plugins.smartfile.fs.fileTreeToObject(this.projectDir, glob)
    );
    
    const smartfilesModArrays = await Promise.all(smartfilesModPromises);
    
    // Flatten the arrays
    const smartfilesMod: plugins.smartfile.SmartFile[] = [];
    smartfilesModArrays.forEach(array => {
      smartfilesMod.push(...array);
    });
    
    // Get test files if not excluded
    let smartfilesTest: plugins.smartfile.SmartFile[] = [];
    if (!excludePaths?.includes('test/')) {
      smartfilesTest = await plugins.smartfile.fs.fileTreeToObject(
        this.projectDir,
        'test/**/*.ts',
      );
    }
    
    return {
      smartfilePackageJSON,
      smartfilesReadme,
      smartfilesReadmeHints,
      smartfilesNpmextraJSON,
      smartfilesMod,
      smartfilesTest,
    };
  }
  
  /**
   * Convert files to context string
   * @param files The files to convert
   * @param mode The context mode to use
   */
  public async convertFilesToContext(
    files: plugins.smartfile.SmartFile[],
    mode: ContextMode = this.contextMode
  ): Promise<string> {
    // Reset context result
    this.contextResult = {
      context: '',
      tokenCount: 0,
      includedFiles: [],
      trimmedFiles: [],
      excludedFiles: [],
      tokenSavings: 0
    };
    
    let totalTokenCount = 0;
    let totalOriginalTokens = 0;

    // Convert SmartFile objects to IFileMetadata for analysis
    const metadata: IFileMetadata[] = files.map(sf => ({
      path: sf.path,
      relativePath: sf.relative,
      size: sf.contents.toString().length,
      mtime: Date.now(), // SmartFile doesn't expose mtime, use current time
      estimatedTokens: this.countTokens(sf.contents.toString()),
      importanceScore: 0
    }));

    // Analyze files using ContextAnalyzer to get smart prioritization
    // (Note: This requires task type which we'll pass from buildContext)
    // For now, sort files by estimated tokens (smaller files first for better efficiency)
    const sortedFiles = [...files].sort((a, b) => {
      const aTokens = this.countTokens(a.contents.toString());
      const bTokens = this.countTokens(b.contents.toString());
      return aTokens - bTokens;
    });

    const processedFiles: string[] = [];

    for (const smartfile of sortedFiles) {
      // Calculate original token count
      const originalContent = smartfile.contents.toString();
      const originalTokenCount = this.countTokens(originalContent);
      totalOriginalTokens += originalTokenCount;
      
      // Apply trimming based on mode
      let processedContent = originalContent;
      
      if (mode !== 'full') {
        processedContent = this.trimmer.trimFile(
          smartfile.relative,
          originalContent,
          mode
        );
      }
      
      // Calculate new token count
      const processedTokenCount = this.countTokens(processedContent);
      
      // Check if we have budget for this file
      if (totalTokenCount + processedTokenCount > this.tokenBudget) {
        // We don't have budget for this file
        this.contextResult.excludedFiles.push({
          path: smartfile.path,
          contents: originalContent,
          relativePath: smartfile.relative,
          tokenCount: originalTokenCount
        });
        continue;
      }
      
      // Format the file for context
      const formattedContent = `
====== START OF FILE ${smartfile.relative} ======
  
${processedContent}
  
====== END OF FILE ${smartfile.relative} ======
      `;
      
      processedFiles.push(formattedContent);
      totalTokenCount += processedTokenCount;
      
      // Track file in appropriate list
      const fileInfo: IFileInfo = {
        path: smartfile.path,
        contents: processedContent,
        relativePath: smartfile.relative,
        tokenCount: processedTokenCount
      };
      
      if (mode === 'full' || processedContent === originalContent) {
        this.contextResult.includedFiles.push(fileInfo);
      } else {
        this.contextResult.trimmedFiles.push(fileInfo);
        this.contextResult.tokenSavings += (originalTokenCount - processedTokenCount);
      }
    }
    
    // Join all processed files
    const context = processedFiles.join('\n');
    
    // Update context result
    this.contextResult.context = context;
    this.contextResult.tokenCount = totalTokenCount;
    
    return context;
  }
  
  /**
   * Convert files to context with smart analysis and prioritization
   * @param metadata - File metadata to analyze
   * @param taskType - Task type for context-aware prioritization
   * @param mode - Context mode to use
   * @returns Context string
   */
  public async convertFilesToContextWithAnalysis(
    metadata: IFileMetadata[],
    taskType: TaskType,
    mode: ContextMode = this.contextMode
  ): Promise<string> {
    // Reset context result
    this.contextResult = {
      context: '',
      tokenCount: 0,
      includedFiles: [],
      trimmedFiles: [],
      excludedFiles: [],
      tokenSavings: 0
    };

    // Analyze files for smart prioritization
    const analysis = await this.analyzer.analyze(metadata, taskType, []);

    // Sort files by importance score (highest first)
    const sortedAnalysis = [...analysis.files].sort(
      (a, b) => b.importanceScore - a.importanceScore
    );

    // Filter out excluded tier
    const relevantFiles = sortedAnalysis.filter(f => f.tier !== 'excluded');

    let totalTokenCount = 0;
    let totalOriginalTokens = 0;
    const processedFiles: string[] = [];

    // Load files with cache support
    for (const fileAnalysis of relevantFiles) {
      try {
        // Check cache first
        let contents: string;
        let originalTokenCount: number;

        const cached = await this.cache.get(fileAnalysis.path);
        if (cached) {
          contents = cached.contents;
          originalTokenCount = cached.tokenCount;
        } else {
          // Load file
          const fileData = await plugins.smartfile.fs.toStringSync(fileAnalysis.path);
          contents = fileData;
          originalTokenCount = this.countTokens(contents);

          // Cache it
          await this.cache.set({
            path: fileAnalysis.path,
            contents,
            tokenCount: originalTokenCount,
            mtime: Date.now(),
            cachedAt: Date.now()
          });
        }

        totalOriginalTokens += originalTokenCount;

        // Apply tier-based trimming
        let processedContent = contents;
        let trimLevel: 'none' | 'light' | 'aggressive' = 'light';

        if (fileAnalysis.tier === 'essential') {
          trimLevel = 'none';
        } else if (fileAnalysis.tier === 'important') {
          trimLevel = 'light';
        } else if (fileAnalysis.tier === 'optional') {
          trimLevel = 'aggressive';
        }

        // Apply trimming based on mode and tier
        if (mode !== 'full' && trimLevel !== 'none') {
          const relativePath = plugins.path.relative(this.projectDir, fileAnalysis.path);
          processedContent = this.trimmer.trimFileWithLevel(
            relativePath,
            contents,
            trimLevel
          );
        }

        // Calculate token count
        const processedTokenCount = this.countTokens(processedContent);

        // Check token budget
        if (totalTokenCount + processedTokenCount > this.tokenBudget) {
          // We don't have budget for this file
          const relativePath = plugins.path.relative(this.projectDir, fileAnalysis.path);
          this.contextResult.excludedFiles.push({
            path: fileAnalysis.path,
            contents,
            relativePath,
            tokenCount: originalTokenCount,
            importanceScore: fileAnalysis.importanceScore
          });
          continue;
        }

        // Format the file for context
        const relativePath = plugins.path.relative(this.projectDir, fileAnalysis.path);
        const formattedContent = `
====== START OF FILE ${relativePath} ======

${processedContent}

====== END OF FILE ${relativePath} ======
        `;

        processedFiles.push(formattedContent);
        totalTokenCount += processedTokenCount;

        // Track file in appropriate list
        const fileInfo: IFileInfo = {
          path: fileAnalysis.path,
          contents: processedContent,
          relativePath,
          tokenCount: processedTokenCount,
          importanceScore: fileAnalysis.importanceScore
        };

        if (trimLevel === 'none' || processedContent === contents) {
          this.contextResult.includedFiles.push(fileInfo);
        } else {
          this.contextResult.trimmedFiles.push(fileInfo);
          this.contextResult.tokenSavings += (originalTokenCount - processedTokenCount);
        }
      } catch (error) {
        console.warn(`Failed to process file ${fileAnalysis.path}:`, error.message);
      }
    }

    // Join all processed files
    const context = processedFiles.join('\n');

    // Update context result
    this.contextResult.context = context;
    this.contextResult.tokenCount = totalTokenCount;

    return context;
  }

  /**
   * Build context for the project
   * @param taskType Optional task type for task-specific context
   */
  public async buildContext(taskType?: TaskType): Promise<IContextResult> {
    // Initialize if needed
    if (this.tokenBudget === 0) {
      await this.initialize();
    }
    
    // Get task-specific configuration if a task type is provided
    if (taskType) {
      const taskConfig = this.configManager.getTaskConfig(taskType);
      if (taskConfig.mode) {
        this.setContextMode(taskConfig.mode);
      }
    }
    
    // Check if analyzer is enabled in config
    const analyzerConfig = this.configManager.getAnalyzerConfig();
    const useAnalyzer = analyzerConfig.enabled && taskType;

    if (useAnalyzer) {
      // Use new smart context building with lazy loading and analysis
      const taskConfig = this.configManager.getTaskConfig(taskType!);

      // Build globs for scanning
      const includeGlobs = taskConfig?.includePaths?.map(p => `${p}/**/*.ts`) || [
        'ts/**/*.ts',
        'ts*/**/*.ts'
      ];

      // Add config files
      const configGlobs = [
        'package.json',
        'readme.md',
        'readme.hints.md',
        'npmextra.json'
      ];

      // Scan files for metadata (fast, doesn't load contents)
      const metadata = await this.lazyLoader.scanFiles([...configGlobs, ...includeGlobs]);

      // Use analyzer to build context with smart prioritization
      await this.convertFilesToContextWithAnalysis(metadata, taskType!, this.contextMode);
    } else {
      // Fall back to old method for backward compatibility
      const taskConfig = taskType ? this.configManager.getTaskConfig(taskType) : undefined;
      const files = await this.gatherFiles(
        taskConfig?.includePaths,
        taskConfig?.excludePaths
      );

      // Convert files to context
      // Create an array of all files to process
      const allFiles: plugins.smartfile.SmartFile[] = [];

      // Add individual files
      if (files.smartfilePackageJSON) allFiles.push(files.smartfilePackageJSON as plugins.smartfile.SmartFile);
      if (files.smartfilesReadme) allFiles.push(files.smartfilesReadme as plugins.smartfile.SmartFile);
      if (files.smartfilesReadmeHints) allFiles.push(files.smartfilesReadmeHints as plugins.smartfile.SmartFile);
      if (files.smartfilesNpmextraJSON) allFiles.push(files.smartfilesNpmextraJSON as plugins.smartfile.SmartFile);

      // Add arrays of files
      if (files.smartfilesMod) {
        if (Array.isArray(files.smartfilesMod)) {
          allFiles.push(...files.smartfilesMod);
        } else {
          allFiles.push(files.smartfilesMod);
        }
      }

      if (files.smartfilesTest) {
        if (Array.isArray(files.smartfilesTest)) {
          allFiles.push(...files.smartfilesTest);
        } else {
          allFiles.push(files.smartfilesTest);
        }
      }

      await this.convertFilesToContext(allFiles);
    }

    return this.contextResult;
  }
  
  /**
   * Update the context with git diff information for commit tasks
   * @param gitDiff The git diff to include
   */
  public updateWithGitDiff(gitDiff: string): IContextResult {
    // If we don't have a context yet, return empty result
    if (!this.contextResult.context) {
      return this.contextResult;
    }
    
    // Add git diff to context
    const diffSection = `
====== GIT DIFF ======

${gitDiff}

====== END GIT DIFF ======
    `;
    
    const diffTokenCount = this.countTokens(diffSection);
    
    // Update context and token count
    this.contextResult.context += diffSection;
    this.contextResult.tokenCount += diffTokenCount;
    
    return this.contextResult;
  }
  
  /**
   * Count tokens in a string
   * @param text The text to count tokens for
   * @param model The model to use for token counting
   */
  public countTokens(text: string, model: string = 'gpt-3.5-turbo'): number {
    try {
      // Use the gpt-tokenizer library to count tokens
      const tokens = plugins.gptTokenizer.encode(text);
      return tokens.length;
    } catch (error) {
      console.error('Error counting tokens:', error);
      // Provide a rough estimate if tokenization fails
      return Math.ceil(text.length / 4);
    }
  }
  
  /**
   * Get the context result
   */
  public getContextResult(): IContextResult {
    return this.contextResult;
  }
  
  /**
   * Get the token count for the current context
   */
  public getTokenCount(): number {
    return this.contextResult.tokenCount;
  }
  
  /**
   * Get both the context string and its token count
   */
  public getContextWithTokenCount(): { context: string; tokenCount: number } {
    return {
      context: this.contextResult.context,
      tokenCount: this.contextResult.tokenCount
    };
  }
}