reports/gemini_token_tracker.py

"""
Gemini API Token Usage Tracker

Tracks token usage across all requests and sends email alerts to admin
when usage reaches configured thresholds (75%, 90%, 100%).

Usage is stored in MongoDB for persistence across server restarts.
"""

import os
import asyncio
from datetime import datetime, timezone
from dotenv import load_dotenv
from reports.database import db
from typing import Optional, Dict, Any
from threading import Lock

load_dotenv()

# Gemini free plan limits:
# - 1,000,000 tokens/month (approximately)
# Default: 1,000,000 tokens (adjust based on your actual plan)
GEMINI_USAGE_LIMIT = int(os.getenv("GEMINI_USAGE_LIMIT", 1000000))
USAGE_ALERT_THRESHOLDS = [75, 90, 100]  # Percentage thresholds for alerts

# In-memory cache to avoid repeated DB reads
_usage_cache: Dict[str, Any] = {
    "total_tokens": 0,
    "last_updated": None,
    "alerts_sent": []  # List of threshold percentages already alerted
}
_cache_lock = Lock()


class GeminiTokenTracker:
    """
    Singleton class for tracking Gemini API token usage.
    """
    
    _instance = None
    _initialized = False
    
    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance
    
    def __init__(self):
        if not self._initialized:
            self._initialized = True
    
    async def initialize(self):
        """Initialize usage data from database."""
        await self._load_from_db()
    
    async def _load_from_db(self):
        """Load current usage from MongoDB."""
        try:
            usage_doc = await db["gemini_usage"].find_one({"_id": "monthly_usage"})
            if usage_doc:
                with _cache_lock:
                    _usage_cache["total_tokens"] = usage_doc.get("total_tokens", 0)
                    _usage_cache["last_updated"] = usage_doc.get("last_updated")
                    _usage_cache["alerts_sent"] = usage_doc.get("alerts_sent", [])
                    _usage_cache["reset_date"] = usage_doc.get("reset_date")
                    _usage_cache["month"] = usage_doc.get("month")
                    _usage_cache["requests_count"] = usage_doc.get("requests_count", 0)

                # Check if we need to reset for new month
                await self._check_monthly_reset()
            else:
                # Initialize new usage document
                await self._reset_usage()
        except Exception as e:
            print(f"⚠️ Error loading usage from DB: {e}")
            # Initialize cache with defaults if DB fails
            with _cache_lock:
                _usage_cache["total_tokens"] = 0
                _usage_cache["requests_count"] = 0
                _usage_cache["alerts_sent"] = []
    
    async def _check_monthly_reset(self):
        """Check if usage should be reset for new month."""
        current_month = datetime.now(timezone.utc).strftime("%Y-%m")
        
        with _cache_lock:
            stored_month = _usage_cache.get("month")
        
        if stored_month != current_month:
            print(f"🔄 New month detected ({current_month}), resetting usage...")
            await self._reset_usage()
    
    async def _reset_usage(self):
        """Reset usage for new month."""
        current_month = datetime.now(timezone.utc).strftime("%Y-%m")
        reset_date = datetime.now(timezone.utc)

        usage_doc = {
            "_id": "monthly_usage",
            "month": current_month,
            "total_tokens": 0,
            "requests_count": 0,
            "alerts_sent": [],
            "reset_date": reset_date,
            "last_updated": reset_date,
            "history": []  # Track daily usage
        }

        await db["gemini_usage"].update_one(
            {"_id": "monthly_usage"},
            {"$set": usage_doc},
            upsert=True
        )

        with _cache_lock:
            _usage_cache["total_tokens"] = 0
            _usage_cache["requests_count"] = 0
            _usage_cache["last_updated"] = reset_date
            _usage_cache["alerts_sent"] = []
            _usage_cache["reset_date"] = reset_date
            _usage_cache["month"] = current_month

        print(f"✅ Usage reset for {current_month}")
    
    async def add_usage(self, tokens: int, endpoint: str = "unknown", user_email: str = None):
        """
        Add token usage to the tracker.

        Args:
            tokens: Number of tokens used
            endpoint: API endpoint that was called
            user_email: Email of user who made the request
        """
        try:
            now = datetime.now(timezone.utc)
            current_month = now.strftime("%Y-%m")
            today = now.strftime("%Y-%m-%d")

            # Update database
            result = await db["gemini_usage"].update_one(
                {"_id": "monthly_usage"},
                {
                    "$inc": {
                        "total_tokens": tokens,
                        "requests_count": 1
                    },
                    "$set": {
                        "last_updated": now,
                        "month": current_month
                    },
                    "$push": {
                        "history": {
                            "date": today,
                            "tokens": tokens,
                            "endpoint": endpoint,
                            "user_email": user_email,
                            "timestamp": now
                        }
                    }
                },
                upsert=True
            )

            # Update cache
            with _cache_lock:
                _usage_cache["total_tokens"] += tokens
                _usage_cache["requests_count"] = _usage_cache.get("requests_count", 0) + 1
                _usage_cache["last_updated"] = now
                _usage_cache["month"] = current_month

            print(f"📊 Token usage updated: +{tokens} tokens | Requests: {_usage_cache['requests_count']} | Total: {_usage_cache['total_tokens']}")

            # Check if we need to send alerts
            await self._check_and_send_alerts()

        except Exception as e:
            print(f"❌ Error adding token usage: {e}")
    
    async def _check_and_send_alerts(self):
        """Check usage percentage and send alerts if thresholds are crossed."""
        with _cache_lock:
            total_tokens = _usage_cache["total_tokens"]
            alerts_sent = set(_usage_cache["alerts_sent"])
        
        usage_percentage = (total_tokens / GEMINI_USAGE_LIMIT) * 100
        
        for threshold in USAGE_ALERT_THRESHOLDS:
            if usage_percentage >= threshold and threshold not in alerts_sent:
                await self._send_admin_alert(threshold, usage_percentage)
                
                # Mark this threshold as alerted
                with _cache_lock:
                    _usage_cache["alerts_sent"].append(threshold)
                
                # Update DB
                await db["gemini_usage"].update_one(
                    {"_id": "monthly_usage"},
                    {"$addToSet": {"alerts_sent": threshold}}
                )
    
    async def _send_admin_alert(self, threshold: int, current_percentage: float):
        """Send email alert to admin via Brevo."""
        admin_email = os.getenv("BREVO_FROM_EMAIL", "subhancontact2@gmail.com")
        if not admin_email:
            print(f"⚠️ BREVO_FROM_EMAIL not configured, cannot send alert for {threshold}% threshold")
            return
        
        try:
            from reports.admin_email_sender import send_token_usage_alert_email
            
            remaining_tokens = GEMINI_USAGE_LIMIT - _usage_cache["total_tokens"]
            remaining_percentage = 100 - current_percentage
            
            result = send_token_usage_alert_email(
                to_email=admin_email,
                threshold=threshold,
                current_percentage=current_percentage,
                total_tokens_used=_usage_cache["total_tokens"],
                total_limit=GEMINI_USAGE_LIMIT,
                remaining_tokens=remaining_tokens,
                remaining_percentage=remaining_percentage
            )
            
            if result["success"]:
                print(f"✅ Admin alert email sent to {admin_email} for {threshold}% threshold")
            else:
                print(f"❌ Failed to send admin alert: {result.get('error')}")
                
        except Exception as e:
            print(f"❌ Error sending admin alert: {e}")
    
    async def get_usage_stats(self) -> Dict[str, Any]:
        """Get current usage statistics - always fetch fresh from DB."""
        # Always reload from DB to ensure fresh data
        await self._load_from_db()

        with _cache_lock:
            total_tokens = _usage_cache["total_tokens"]
            alerts_sent = list(_usage_cache.get("alerts_sent", []))
            reset_date = _usage_cache.get("reset_date")
            month = _usage_cache.get("month")

        usage_percentage = (total_tokens / GEMINI_USAGE_LIMIT) * 100 if GEMINI_USAGE_LIMIT > 0 else 0
        remaining_tokens = GEMINI_USAGE_LIMIT - total_tokens

        stats = {
            "month": month,
            "total_tokens_used": total_tokens,
            "total_limit": GEMINI_USAGE_LIMIT,
            "usage_percentage": round(usage_percentage, 2),
            "remaining_tokens": remaining_tokens,
            "remaining_percentage": round(100 - usage_percentage, 2),
            "alerts_sent": alerts_sent,
            "reset_date": reset_date.isoformat() if reset_date else None,
            "last_updated": _usage_cache["last_updated"].isoformat() if _usage_cache["last_updated"] else None,
            "requests_count": _usage_cache.get("requests_count", 0)
        }
        
        print(f"📊 Gemini Usage Stats: {total_tokens:,} / {GEMINI_USAGE_LIMIT:,} tokens ({usage_percentage:.2f}%)")
        return stats


# Global tracker instance
_tracker: Optional[GeminiTokenTracker] = None


def get_token_tracker() -> GeminiTokenTracker:
    """Get the global token tracker instance."""
    global _tracker
    if _tracker is None:
        _tracker = GeminiTokenTracker()
    return _tracker


async def initialize_token_tracker():
    """Initialize the token tracker (call this on app startup)."""
    tracker = get_token_tracker()
    await tracker.initialize()
    print(f"✅ Gemini Token Tracker initialized (Limit: {GEMINI_USAGE_LIMIT:,} tokens/month)")


async def track_token_usage(tokens: int, endpoint: str = "audit", user_email: str = None):
    """
    Convenience function to track token usage.

    Args:
        tokens: Number of tokens used
        endpoint: API endpoint name
        user_email: User's email address
    """
    tracker = get_token_tracker()
    await tracker.add_usage(tokens, endpoint, user_email)


async def track_api_error(endpoint: str = "unknown", error_type: str = "unknown", user_email: str = None):
    """
    Track API errors (like rate limits) without consuming tokens.
    Useful for monitoring failed requests.
    
    Args:
        endpoint: API endpoint that failed
        error_type: Type of error (e.g., "rate_limit", "auth_error")
        user_email: User's email address
    """
    tracker = get_token_tracker()
    # Track as -1 tokens to indicate error (won't affect usage stats)
    await tracker.add_usage(0, f"{endpoint}_error_{error_type}", user_email)
    print(f"🚨 API error tracked: {endpoint} - {error_type}")


async def get_current_usage() -> Dict[str, Any]:
    """Get current usage statistics."""
    tracker = get_token_tracker()
    return await tracker.get_usage_stats()