import { Injectable, Logger } from '@nestjs/common';

export interface MemoryStats {
  heapUsed: number;    
  heapTotal: number;   
  external: number;    
  rss: number;         
  timestamp: Date;
}

@Injectable()
export class MemoryMonitorService {
  private readonly logger = new Logger(MemoryMonitorService.name);
  private readonly MAX_MEMORY_MB: number;
  private readonly BATCH_SIZE: number;
  private readonly GC_THRESHOLD_MB: number;

  constructor() {
    
    this.MAX_MEMORY_MB = parseInt(process.env.MAX_MEMORY_USAGE_MB || '1024'); 
    this.BATCH_SIZE = parseInt(process.env.CHUNK_BATCH_SIZE || '100'); 
    this.GC_THRESHOLD_MB = parseInt(process.env.GC_THRESHOLD_MB || '800'); 

    this.logger.log(`Initialized memory monitoring: Limit=${this.MAX_MEMORY_MB}MB, Batch Size=${this.BATCH_SIZE}, GC Threshold=${this.GC_THRESHOLD_MB}MB`);
  }

  
  getMemoryUsage(): MemoryStats {
    const usage = process.memoryUsage();
    return {
      heapUsed: Math.round(usage.heapUsed / 1024 / 1024),
      heapTotal: Math.round(usage.heapTotal / 1024 / 1024),
      external: Math.round((usage.external || 0) / 1024 / 1024),
      rss: Math.round(usage.rss / 1024 / 1024),
      timestamp: new Date(),
    };
  }

  
  isMemoryHigh(): boolean {
    const usage = this.getMemoryUsage();
    return usage.heapUsed > this.MAX_MEMORY_MB * 0.85; 
  }

  
  async waitForMemoryAvailable(timeoutMs: number = 30000): Promise<void> {
    const startTime = Date.now();

    while (this.isMemoryHigh()) {
      if (Date.now() - startTime > timeoutMs) {
        throw new Error(`Memory wait timed out: Currently ${this.getMemoryUsage().heapUsed}MB > ${this.MAX_MEMORY_MB * 0.85}MB`);
      }

      this.logger.warn(
        `Memory usage is too high. Waiting for release... ${this.getMemoryUsage().heapUsed}/${this.MAX_MEMORY_MB}MB`,
      );

      
      if (global.gc) {
        this.logger.log('Running forced garbage collection...');
        global.gc();
      }

      await new Promise((resolve) => setTimeout(resolve, 1000));
    }
  }

  
  forceGC(): void {
    if (global.gc) {
      const before = this.getMemoryUsage();
      global.gc();
      const after = this.getMemoryUsage();
      this.logger.log(
        `GC Complete: ${before.heapUsed}MB → ${after.heapUsed}MB (${before.heapUsed - after.heapUsed}MB released)`,
      );
    }
  }

  
  getDynamicBatchSize(currentMemoryMB: number): number {
    const baseBatchSize = this.BATCH_SIZE;

    if (currentMemoryMB > this.GC_THRESHOLD_MB) {
      
      const reduced = Math.max(10, Math.floor(baseBatchSize * 0.5));
      this.logger.warn(
        `Memory constrained (${currentMemoryMB}MB), dynamically adjusting batch size: ${baseBatchSize} → ${reduced}`,
      );
      return reduced;
    } else if (currentMemoryMB < this.MAX_MEMORY_MB * 0.4) {
      
      const increased = Math.min(200, Math.floor(baseBatchSize * 1.2));
      if (increased > baseBatchSize) {
        this.logger.log(
          `Memory available (${currentMemoryMB}MB), dynamically adjusting batch size: ${baseBatchSize} → ${increased}`,
        );
      }
      return increased;
    }

    return baseBatchSize;
  }

  
  async processInBatches<T, R>(
    items: T[],
    processor: (batch: T[], batchIndex: number) => Promise<R[]>,
    options?: {
      onBatchComplete?: (batchIndex: number, totalBatches: number, results: R[]) => Promise<void> | void;
      maxConcurrency?: number;
    },
  ): Promise<R[]> {
    const totalItems = items.length;
    if (totalItems === 0) return [];

    const startTime = Date.now();
    this.logger.log(`Starting batch processing: ${totalItems} items`);

    const allResults: R[] = [];
    let processedCount = 0;

    for (let i = 0; i < totalItems;) {
      
      await this.waitForMemoryAvailable();

      
      const currentMem = this.getMemoryUsage().heapUsed;
      const batchSize = this.getDynamicBatchSize(currentMem);

      
      const batch = items.slice(i, i + batchSize);
      const batchIndex = Math.floor(i / batchSize) + 1;
      const totalBatches = Math.ceil(totalItems / batchSize);

      this.logger.log(
        `Processing batch ${batchIndex}/${totalBatches}: ${batch.length} items (Total ${processedCount}/${totalItems})`,
      );

      
      const batchResults = await processor(batch, batchIndex);
      allResults.push(...batchResults);
      processedCount += batch.length;

      
      if (options?.onBatchComplete) {
        await options.onBatchComplete(batchIndex, totalBatches, batchResults);
      }

      
      if (currentMem > this.GC_THRESHOLD_MB) {
        this.forceGC();
      }

      
      batch.length = 0;

      i += batchSize;
    }

    const duration = ((Date.now() - startTime) / 1000).toFixed(2);
    const finalMem = this.getMemoryUsage();
    this.logger.log(
      `Batch processing complete: ${totalItems} items, Duration ${duration}s, Final memory ${finalMem.heapUsed}MB`,
    );

    return allResults;
  }

  
  estimateMemoryUsage(itemCount: number, itemSizeBytes: number, vectorDim: number): number {
    
    const textMemory = itemCount * itemSizeBytes;

    
    const vectorMemory = itemCount * vectorDim * 4;

    
    const overhead = itemCount * 100;

    const totalMB = Math.round((textMemory + vectorMemory + overhead) / 1024 / 1024);

    return totalMB;
  }

  
  shouldUseBatching(itemCount: number, itemSizeBytes: number, vectorDim: number): boolean {
    const estimatedMB = this.estimateMemoryUsage(itemCount, itemSizeBytes, vectorDim);
    const threshold = this.MAX_MEMORY_MB * 0.7; 

    if (estimatedMB > threshold) {
      this.logger.warn(
        `Estimated memory ${estimatedMB}MB exceeds threshold ${threshold}MB, using batch processing`,
      );
      return true;
    }

    return false;
  }

  
  getRecommendedBatchSize(itemSizeBytes: number, vectorDim: number): number {
    
    const targetMemoryMB = 200;
    const targetMemoryBytes = targetMemoryMB * 1024 * 1024;

    
    const singleItemMemory = itemSizeBytes + (vectorDim * 4) + 100;

    const batchSize = Math.floor(targetMemoryBytes / singleItemMemory);

    
    return Math.max(10, Math.min(200, batchSize));
  }
}
