KnowledgeBase/server/src/vision-pipeline/cost-control.service.ts

import { Injectable, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { InjectRepository } from '@nestjs/typeorm';
import { Repository } from 'typeorm';
import { User } from '../user/user.entity';

export interface UserQuota {
  userId: string;
  monthlyCost: number;      
  maxCost: number;          
  remaining: number;        
  lastReset: Date;          
}

export interface CostEstimate {
  estimatedCost: number;    
  estimatedTime: number;    
  pageBreakdown: {          
    pageIndex: number;
    cost: number;
  }[];
}

@Injectable()
export class CostControlService {
  private readonly logger = new Logger(CostControlService.name);
  private readonly COST_PER_PAGE = 0.01; 
  private readonly DEFAULT_MONTHLY_LIMIT = 100; 

  constructor(
    private configService: ConfigService,
    @InjectRepository(User)
    private userRepository: Repository<User>,
  ) { }

  
  estimateCost(pageCount: number, quality: 'low' | 'medium' | 'high' = 'medium'): CostEstimate {
    
    const qualityMultiplier = {
      low: 0.5,
      medium: 1.0,
      high: 1.5,
    };

    const baseCost = pageCount * this.COST_PER_PAGE * qualityMultiplier[quality];
    const estimatedTime = pageCount * 3; 

    const pageBreakdown = Array.from({ length: pageCount }, (_, i) => ({
      pageIndex: i + 1,
      cost: this.COST_PER_PAGE * qualityMultiplier[quality],
    }));

    return {
      estimatedCost: baseCost,
      estimatedTime,
      pageBreakdown,
    };
  }

  
  async checkQuota(userId: string, estimatedCost: number): Promise<{
    allowed: boolean;
    quota: UserQuota;
    reason?: string;
  }> {
    const quota = await this.getUserQuota(userId);

    
    this.checkAndResetMonthlyQuota(quota);

    if (quota.remaining < estimatedCost) {
      this.logger.warn(
        `Insufficient quota for user ${userId}: Remaining $${quota.remaining.toFixed(2)}, Required $${estimatedCost.toFixed(2)}`,
      );
      return {
        allowed: false,
        quota,
        reason: `Insufficient quota: remaining $${quota.remaining.toFixed(2)}, required $${estimatedCost.toFixed(2)}`,
      };
    }

    return {
      allowed: true,
      quota,
    };
  }

  
  async deductQuota(userId: string, actualCost: number): Promise<void> {
    const quota = await this.getUserQuota(userId);
    quota.monthlyCost += actualCost;
    quota.remaining = quota.maxCost - quota.monthlyCost;

    await this.userRepository.update(userId, {
      monthlyCost: quota.monthlyCost,
    });

    this.logger.log(
      `Deducted $${actualCost.toFixed(2)} from user ${userId}'s quota. Remaining $${quota.remaining.toFixed(2)}`,
    );
  }

  
  async getUserQuota(userId: string): Promise<UserQuota> {
    const user = await this.userRepository.findOne({ where: { id: userId } });

    if (!user) {
      throw new Error(`User ${userId} does not exist`);
    }

    
    const monthlyCost = user.monthlyCost || 0;
    const maxCost = user.maxCost || this.DEFAULT_MONTHLY_LIMIT;
    const lastReset = user.lastQuotaReset || new Date();

    return {
      userId,
      monthlyCost,
      maxCost,
      remaining: maxCost - monthlyCost,
      lastReset,
    };
  }

  
  private checkAndResetMonthlyQuota(quota: UserQuota): void {
    const now = new Date();
    const lastReset = quota.lastReset;

    
    if (
      now.getMonth() !== lastReset.getMonth() ||
      now.getFullYear() !== lastReset.getFullYear()
    ) {
      this.logger.log(`Reset monthly quota for user ${quota.userId}`);

      
      quota.monthlyCost = 0;
      quota.remaining = quota.maxCost;
      quota.lastReset = now;

      
      this.userRepository.update(quota.userId, {
        monthlyCost: 0,
        lastQuotaReset: now,
      });
    }
  }

  
  async setQuotaLimit(userId: string, maxCost: number): Promise<void> {
    await this.userRepository.update(userId, { maxCost });
    this.logger.log(`Set quota limit for user ${userId} to $${maxCost}`);
  }

  
  async getCostReport(userId: string, days: number = 30): Promise<{
    totalCost: number;
    dailyAverage: number;
    pageStats: {
      totalPages: number;
      avgCostPerPage: number;
    };
    quotaUsage: number; 
  }> {
    const quota = await this.getUserQuota(userId);
    const usagePercent = (quota.monthlyCost / quota.maxCost) * 100;

    
    

    return {
      totalCost: quota.monthlyCost,
      dailyAverage: quota.monthlyCost / Math.max(days, 1),
      pageStats: {
        totalPages: Math.floor(quota.monthlyCost / this.COST_PER_PAGE),
        avgCostPerPage: this.COST_PER_PAGE,
      },
      quotaUsage: usagePercent,
    };
  }

  
  async checkWarningThreshold(userId: string): Promise<{
    shouldWarn: boolean;
    message: string;
  }> {
    const quota = await this.getUserQuota(userId);
    const usagePercent = (quota.monthlyCost / quota.maxCost) * 100;

    if (usagePercent >= 90) {
      return {
        shouldWarn: true,
        message: `⚠️ Quota usage has reached ${usagePercent.toFixed(1)}%. Remaining $${quota.remaining.toFixed(2)}`,
      };
    }

    if (usagePercent >= 75) {
      return {
        shouldWarn: true,
        message: `💡 Quota usage ${usagePercent.toFixed(1)}%. Be careful with controlling costs`,
      };
    }

    return {
      shouldWarn: false,
      message: '',
    };
  }

  
  formatCost(cost: number): string {
    return `$${cost.toFixed(2)}`;
  }

  
  formatTime(seconds: number): string {
    if (seconds < 60) {
      return `${seconds.toFixed(0)}s`;
    }
    const minutes = Math.floor(seconds / 60);
    const remainingSeconds = seconds % 60;
    return `${minutes}m ${remainingSeconds.toFixed(0)}s`;
  }
}