app/utils/token_tracker.py

"""
Token usage tracking for cost analysis and monitoring.

Tracks:
- Input/output tokens per request
- Cost per request
- Cumulative token usage
- Feature-specific usage
"""

import logging
from datetime import datetime
from typing import Dict, List, Optional
from dataclasses import dataclass, asdict
from threading import Lock

logger = logging.getLogger(__name__)


@dataclass
class TokenRecord:
    """Record of a single token usage event."""
    timestamp: datetime
    feature: str  # 'text_triage', 'image_scan', 'video_analysis'
    input_tokens: int
    output_tokens: int
    total_tokens: int
    cost_usd: float
    owner_id: str
    model: str  # 'gemini-1.5-pro', etc.
    status: str  # 'success', 'error', 'rate_limited'
    latency_ms: float


class TokenTracker:
    """Track token usage for cost analysis and monitoring."""
    
    def __init__(self):
        self.records: List[TokenRecord] = []
        self.lock = Lock()
        
        # Pricing (Gemini 1.5 Flash - update as needed)
        # https://ai.google.dev/pricing
        self.pricing = {
            "gemini-1.5-flash": {
                "input_cost_per_million": 0.075,
                "output_cost_per_million": 0.30
            },
            "gemini-1.5-pro": {
                "input_cost_per_million": 1.50,
                "output_cost_per_million": 6.00
            },
            "gemini-2.0-flash": {
                "input_cost_per_million": 0.10,
                "output_cost_per_million": 0.40
            }
        }
    
    def calculate_cost(self, model: str, input_tokens: int, output_tokens: int) -> float:
        """Calculate cost in USD for token usage."""
        if model not in self.pricing:
            logger.warning(f"Unknown model: {model}, using gemini-1.5-flash pricing")
            model = "gemini-1.5-flash"
        
        pricing = self.pricing[model]
        input_cost = (input_tokens / 1_000_000) * pricing["input_cost_per_million"]
        output_cost = (output_tokens / 1_000_000) * pricing["output_cost_per_million"]
        
        return input_cost + output_cost
    
    def log_tokens(
        self,
        feature: str,
        input_tokens: int,
        output_tokens: int,
        owner_id: str,
        model: str = "gemini-1.5-flash",
        status: str = "success",
        latency_ms: float = 0.0
    ) -> TokenRecord:
        """Log a token usage event."""
        total_tokens = input_tokens + output_tokens
        cost = self.calculate_cost(model, input_tokens, output_tokens)
        
        record = TokenRecord(
            timestamp=datetime.utcnow(),
            feature=feature,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=total_tokens,
            cost_usd=cost,
            owner_id=owner_id,
            model=model,
            status=status,
            latency_ms=latency_ms
        )
        
        with self.lock:
            self.records.append(record)
        
        # Log with structured logging
        logger.info(
            "Token usage recorded",
            extra={
                "feature": feature,
                "input_tokens": input_tokens,
                "output_tokens": output_tokens,
                "total_tokens": total_tokens,
                "cost_usd": round(cost, 4),
                "owner_id": owner_id,
                "model": model,
                "status": status,
                "latency_ms": latency_ms
            }
        )
        
        return record
    
    def get_stats(self, owner_id: Optional[str] = None) -> Dict:
        """Get token usage statistics."""
        with self.lock:
            if owner_id:
                records = [r for r in self.records if r.owner_id == owner_id]
            else:
                records = self.records
        
        if not records:
            return {
                "total_requests": 0,
                "total_tokens": 0,
                "total_cost_usd": 0.0,
                "average_tokens_per_request": 0,
                "average_cost_per_request": 0.0
            }
        
        total_requests = len(records)
        total_tokens = sum(r.total_tokens for r in records)
        total_cost = sum(r.cost_usd for r in records)
        
        return {
            "total_requests": total_requests,
            "total_tokens": total_tokens,
            "total_cost_usd": round(total_cost, 4),
            "average_tokens_per_request": total_tokens // total_requests if total_requests else 0,
            "average_cost_per_request": round(total_cost / total_requests, 4) if total_requests else 0.0,
            "by_feature": self._get_feature_breakdown(records),
            "by_status": self._get_status_breakdown(records)
        }
    
    def _get_feature_breakdown(self, records: List[TokenRecord]) -> Dict:
        """Get token usage by feature."""
        breakdown = {}
        for record in records:
            if record.feature not in breakdown:
                breakdown[record.feature] = {
                    "requests": 0,
                    "tokens": 0,
                    "cost_usd": 0.0
                }
            breakdown[record.feature]["requests"] += 1
            breakdown[record.feature]["tokens"] += record.total_tokens
            breakdown[record.feature]["cost_usd"] += record.cost_usd
        
        # Round costs
        for feature in breakdown:
            breakdown[feature]["cost_usd"] = round(breakdown[feature]["cost_usd"], 4)
        
        return breakdown
    
    def _get_status_breakdown(self, records: List[TokenRecord]) -> Dict:
        """Get token usage by status."""
        breakdown = {}
        for record in records:
            if record.status not in breakdown:
                breakdown[record.status] = {
                    "count": 0,
                    "total_tokens": 0,
                    "total_cost_usd": 0.0
                }
            breakdown[record.status]["count"] += 1
            breakdown[record.status]["total_tokens"] += record.total_tokens
            breakdown[record.status]["total_cost_usd"] += record.cost_usd
        
        # Round costs
        for status in breakdown:
            breakdown[status]["total_cost_usd"] = round(breakdown[status]["total_cost_usd"], 4)
        
        return breakdown
    
    def get_recent_records(self, limit: int = 100) -> List[Dict]:
        """Get recent token usage records."""
        with self.lock:
            recent = sorted(self.records, key=lambda r: r.timestamp, reverse=True)[:limit]
        
        return [asdict(r) for r in recent]
    
    def export_csv(self, filepath: str, owner_id: Optional[str] = None):
        """Export token records to CSV for analysis."""
        import csv
        
        with self.lock:
            if owner_id:
                records = [r for r in self.records if r.owner_id == owner_id]
            else:
                records = self.records
        
        if not records:
            logger.warning(f"No records to export for {owner_id or 'all users'}")
            return
        
        with open(filepath, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=TokenRecord.__dataclass_fields__.keys())
            writer.writeheader()
            for record in records:
                writer.writerow(asdict(record))
        
        logger.info(f"Exported {len(records)} token records to {filepath}")


# Global instance
token_tracker = TokenTracker()