import { EventEmitter } from 'events';
import { Logger } from './logger.js';

export interface CostUsage {
  date: string;
  totalCost: number;
  apiCalls: number;
  tokens: {
    prompt: number;
    completion: number;
    total: number;
  };
  models: Record<string, {
    calls: number;
    cost: number;
    tokens: number;
  }>;
}

export interface CostTrackerConfig {
  dailyLimit: number;
  monthlyLimit: number;
  warningThreshold: number;
}

export interface TokenUsage {
  promptTokens: number;
  completionTokens: number;
  totalTokens: number;
}

export class CostTracker extends EventEmitter {
  private logger: Logger;
  private config: CostTrackerConfig;
  private dailyUsage: Map<string, CostUsage>;
  private monthlyUsage: Map<string, CostUsage>;
  
  // OpenAI pricing (as of 2025, adjust as needed)
  private readonly modelPricing = {
    'gpt-5': {
      prompt: 0.05 / 1000, // $0.05 per 1K prompt tokens (estimated)
      completion: 0.15 / 1000, // $0.15 per 1K completion tokens (estimated)
    },
    'gpt-4': {
      prompt: 0.03 / 1000, // $0.03 per 1K prompt tokens
      completion: 0.06 / 1000, // $0.06 per 1K completion tokens
    },
    'gpt-4-turbo': {
      prompt: 0.01 / 1000,
      completion: 0.03 / 1000,
    },
    'gpt-3.5-turbo': {
      prompt: 0.0005 / 1000,
      completion: 0.0015 / 1000,
    },
  };

  constructor(config: CostTrackerConfig) {
    super();
    this.config = config;
    this.logger = new Logger('CostTracker');
    this.dailyUsage = new Map();
    this.monthlyUsage = new Map();
    
    this.logger.info('Cost tracker initialized', {
      dailyLimit: config.dailyLimit,
      monthlyLimit: config.monthlyLimit,
      warningThreshold: config.warningThreshold,
    });
  }

  public async trackUsage(
    model: string,
    tokenUsage: TokenUsage,
    cost?: number
  ): Promise<void> {
    const now = new Date();
    const dateKey = now.toISOString().split('T')[0]; // YYYY-MM-DD
    const monthKey = dateKey.substring(0, 7); // YYYY-MM

    // Calculate cost if not provided
    const calculatedCost = cost || this.calculateCost(model, tokenUsage);

    // Update daily usage
    await this.updateUsage(this.dailyUsage, dateKey, model, tokenUsage, calculatedCost);

    // Update monthly usage
    await this.updateUsage(this.monthlyUsage, monthKey, model, tokenUsage, calculatedCost);

    // Check limits and emit warnings
    await this.checkLimits(dateKey, monthKey);

    this.logger.debug('Usage tracked', {
      model,
      cost: calculatedCost,
      tokens: tokenUsage,
      date: dateKey,
      month: monthKey,
    });
  }

  private calculateCost(model: string, tokenUsage: TokenUsage): number {
    const pricing = this.modelPricing[model as keyof typeof this.modelPricing];
    
    if (!pricing) {
      this.logger.warn(`Unknown model pricing for ${model}, using GPT-5 rates`);
      return this.calculateCost('gpt-5', tokenUsage);
    }

    const promptCost = tokenUsage.promptTokens * pricing.prompt;
    const completionCost = tokenUsage.completionTokens * pricing.completion;
    
    return promptCost + completionCost;
  }

  private async updateUsage(
    usageMap: Map<string, CostUsage>,
    key: string,
    model: string,
    tokenUsage: TokenUsage,
    cost: number
  ): Promise<void> {
    let usage = usageMap.get(key);
    
    if (!usage) {
      usage = {
        date: key,
        totalCost: 0,
        apiCalls: 0,
        tokens: {
          prompt: 0,
          completion: 0,
          total: 0,
        },
        models: {},
      };
      usageMap.set(key, usage);
    }

    // Update totals
    usage.totalCost += cost;
    usage.apiCalls += 1;
    usage.tokens.prompt += tokenUsage.promptTokens;
    usage.tokens.completion += tokenUsage.completionTokens;
    usage.tokens.total += tokenUsage.totalTokens;

    // Update model-specific usage
    if (!usage.models[model]) {
      usage.models[model] = {
        calls: 0,
        cost: 0,
        tokens: 0,
      };
    }

    usage.models[model].calls += 1;
    usage.models[model].cost += cost;
    usage.models[model].tokens += tokenUsage.totalTokens;
  }

  private async checkLimits(dateKey: string, monthKey: string): Promise<void> {
    const dailyUsage = this.dailyUsage.get(dateKey);
    const monthlyUsage = this.monthlyUsage.get(monthKey);

    if (dailyUsage) {
      const dailyPercent = dailyUsage.totalCost / this.config.dailyLimit;
      
      if (dailyPercent >= 1.0) {
        this.emit('limit_exceeded', {
          type: 'daily',
          usage: dailyUsage,
          limit: this.config.dailyLimit,
        });
      } else if (dailyPercent >= this.config.warningThreshold) {
        this.emit('warning', {
          type: 'daily',
          usage: dailyUsage,
          limit: this.config.dailyLimit,
          threshold: this.config.warningThreshold,
        });
      }
    }

    if (monthlyUsage) {
      const monthlyPercent = monthlyUsage.totalCost / this.config.monthlyLimit;
      
      if (monthlyPercent >= 1.0) {
        this.emit('limit_exceeded', {
          type: 'monthly',
          usage: monthlyUsage,
          limit: this.config.monthlyLimit,
        });
      } else if (monthlyPercent >= this.config.warningThreshold) {
        this.emit('warning', {
          type: 'monthly',
          usage: monthlyUsage,
          limit: this.config.monthlyLimit,
          threshold: this.config.warningThreshold,
        });
      }
    }
  }

  public getDailyUsage(date?: string): CostUsage | undefined {
    const key = date || new Date().toISOString().split('T')[0];
    return this.dailyUsage.get(key);
  }

  public getMonthlyUsage(month?: string): CostUsage | undefined {
    const key = month || new Date().toISOString().substring(0, 7);
    return this.monthlyUsage.get(key);
  }

  public isWithinDailyLimit(additionalCost: number = 0): boolean {
    const today = new Date().toISOString().split('T')[0];
    const usage = this.dailyUsage.get(today);
    const currentCost = usage?.totalCost || 0;
    return (currentCost + additionalCost) <= this.config.dailyLimit;
  }

  public isWithinMonthlyLimit(additionalCost: number = 0): boolean {
    const thisMonth = new Date().toISOString().substring(0, 7);
    const usage = this.monthlyUsage.get(thisMonth);
    const currentCost = usage?.totalCost || 0;
    return (currentCost + additionalCost) <= this.config.monthlyLimit;
  }

  public getRemainingDailyBudget(): number {
    const today = new Date().toISOString().split('T')[0];
    const usage = this.dailyUsage.get(today);
    const currentCost = usage?.totalCost || 0;
    return Math.max(0, this.config.dailyLimit - currentCost);
  }

  public getRemainingMonthlyBudget(): number {
    const thisMonth = new Date().toISOString().substring(0, 7);
    const usage = this.monthlyUsage.get(thisMonth);
    const currentCost = usage?.totalCost || 0;
    return Math.max(0, this.config.monthlyLimit - currentCost);
  }

  public getUsageSummary(): {
    daily: CostUsage | undefined;
    monthly: CostUsage | undefined;
    limits: CostTrackerConfig;
  } {
    const today = new Date().toISOString().split('T')[0];
    const thisMonth = new Date().toISOString().substring(0, 7);
    
    return {
      daily: this.getDailyUsage(today),
      monthly: this.getMonthlyUsage(thisMonth),
      limits: { ...this.config },
    };
  }
}