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

/**
 * Cost control and quota management service
 * Used to manage API call costs for Vision Pipeline
 */

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;      // Current month used cost
  maxCost: number;          // Monthly max cost
  remaining: number;        // Remaining cost
  lastReset: Date;          // Last reset time
}

export interface CostEstimate {
  estimatedCost: number;    // Estimated cost
  estimatedTime: number;    // Estimated time(seconds)
  pageBreakdown: {          // Per-page breakdown
    pageIndex: number;
    cost: number;
  }[];
}

@Injectable()
export class CostControlService {
  private readonly logger = new Logger(CostControlService.name);
  private readonly COST_PER_PAGE = 0.01; // Cost per page(USD)
  private readonly DEFAULT_MONTHLY_LIMIT = 100; // Default monthly limit(USD)

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

  /**
   * Estimate processing cost
   */
  estimateCost(pageCount: number, quality: 'low' | 'medium' | 'high' = 'medium'): CostEstimate {
    // Adjust cost coefficient based on quality
    const qualityMultiplier = {
      low: 0.5,
      medium: 1.0,
      high: 1.5,
    };

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

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

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

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

    // Check monthly reset
    this.checkAndResetMonthlyQuota(quota);

    if (quota.remaining < estimatedCost) {
      this.logger.warn(
        `User ${userId} quota insufficient: 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,
    };
  }

  /**
   * Deduct from 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} quota. Remaining: $${quota.remaining.toFixed(2)}`,
    );
  }

  /**
   * Get user quota
   */
  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`);
    }

    // Use default if user has no quota info
    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,
    };
  }

  /**
   * Check and reset monthly quota
   */
  private checkAndResetMonthlyQuota(quota: UserQuota): void {
    const now = new Date();
    const lastReset = quota.lastReset;

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

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

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

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

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

    // Query history records here(if implemented)
    // Return current quota info temporarily

    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,
    };
  }

  /**
   * Check cost warning threshold
   */
  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 reached ${usagePercent.toFixed(1)}%. Remaining: $${quota.remaining.toFixed(2)}`,
      };
    }

    if (usagePercent >= 75) {
      return {
        shouldWarn: true,
        message: `💡 Quota usage at ${usagePercent.toFixed(1)}%. Please monitor your costs carefully`,
      };
    }

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

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

  /**
   * Format time display
   */
  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`;
  }
}