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, ) { } 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 { 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 { 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 { 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`; } }