config/global_state.py

import asyncio
import logging
import threading
import time
from collections import defaultdict
from typing import Dict, Optional, Set

from config.settings import MODEL_QUOTA_LIMITS, QUOTA_HARD_LIMIT, QUOTA_SOFT_LIMIT
from models.exceptions import QuotaExceededError

logger = logging.getLogger(__name__)


class GlobalState:
    """
    Singleton class to hold global application state, specifically for Quota Exceeded logic.
    """

    _instance = None
    IS_QUOTA_EXCEEDED = False
    NEEDS_ROTATION = False  # [GR-01] Soft Signal Flag
    QUOTA_EXCEEDED_TIMESTAMP = 0.0

    # Global Event for holding requests during auth rotation
    # Initially set to True (allowed) by init_rotation_lock()
    AUTH_ROTATION_LOCK = asyncio.Event()

    # Global Event to signal Quota Exceeded immediately
    QUOTA_EXCEEDED_EVENT = asyncio.Event()

    # Event to signal that a rotation operation has completed.
    rotation_complete_event = asyncio.Event()

    # Track the type of the last error for adaptive cooldowns
    # Values: 'RATE_LIMIT', 'QUOTA_EXCEEDED', or None
    last_error_type = None

    # Token usage tracking for proactive rotation
    # [QUOTA-02] Changed from single int to dict for model-specific tracking
    current_profile_model_usage: Dict[str, int] = defaultdict(int)
    current_profile_exhausted_models: Set[str] = set()

    # [FINAL-02] Dynamic Rotation Guard: Track queued requests
    queued_request_count = 0

    # Global Shutdown Event (Thread-safe for signal handlers)
    # Used to circuit-break logic during aggressive shutdown
    IS_SHUTTING_DOWN = threading.Event()

    # [ID-01] Global Recovery State Manager
    # Flag to indicate if a recovery operation (auth rotation) is currently active
    IS_RECOVERING = False
    # Flag to indicate if the system is in emergency operation mode (e.g. all profiles exhausted)
    DEPLOYMENT_EMERGENCY_MODE = False
    # Event to signal that recovery is complete and streams can resume
    RECOVERY_EVENT = asyncio.Event()
    # [FIX-RACE] Track last rotation timestamp to handle race conditions
    LAST_ROTATION_TIMESTAMP = 0.0

    # [CONCURRENCY-FIX] Track the currently active stream request ID
    # Used to gate consumers and prevent zombie streams from processing data
    CURRENT_STREAM_REQ_ID: Optional[str] = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(GlobalState, cls).__new__(cls)
        return cls._instance

    @classmethod
    def start_recovery(cls):
        """
        Signals the start of a recovery operation (e.g., auth rotation).
        Pauses streams and sets the recovery flag.
        """
        cls.IS_RECOVERING = True
        cls.RECOVERY_EVENT.clear()
        logger.info("🔄 SYSTEM: Recovery Mode Initiated. Streams will pause.")

    @classmethod
    def finish_recovery(cls):
        """
        Signals the completion of a recovery operation.
        Resumes streams and clears the recovery flag.
        """
        cls.IS_RECOVERING = False
        cls.RECOVERY_EVENT.set()
        # [FIX-RACE] Update timestamp on finish
        cls.LAST_ROTATION_TIMESTAMP = time.time()
        logger.info(
            f"✅ SYSTEM: Recovery Mode Finished at {cls.LAST_ROTATION_TIMESTAMP}. Streams resuming."
        )

    @classmethod
    def init_rotation_lock(cls):
        """Initialize the rotation lock to allow requests."""
        cls.AUTH_ROTATION_LOCK.set()
        logger.info("🔐 Global Auth Rotation Lock initialized (OPEN).")

    @classmethod
    def set_quota_exceeded(cls, message: str = "", model_id: Optional[str] = None):
        """
        Sets the global quota exceeded flag and logs a critical warning.
        Also determines the error type based on the message for adaptive cooldowns.
        Optionally accepts a model_id to flag specific model exhaustion.
        """
        if not cls.IS_QUOTA_EXCEEDED:
            cls.IS_QUOTA_EXCEEDED = True
            cls.QUOTA_EXCEEDED_TIMESTAMP = time.time()
            cls.QUOTA_EXCEEDED_EVENT.set()

            # Determine error type
            safe_message = message if message else ""
            msg_lower = safe_message.lower()
            if (
                "429" in msg_lower
                or "rate limit" in msg_lower
                or "resource has been exhausted" in msg_lower
            ):
                # API "RESOURCE_EXHAUSTED" usually means 429/quota shared behavior,
                # but "rate limit" specifically implies a temporary 429.
                # However, Gemini "Resource has been exhausted" is often a harder limit.
                # Let's verify standard Gemini strings:
                # "429: Too Many Requests" -> Rate Limit
                # "429: Resource has been exhausted" -> Quota
                if "too many requests" in msg_lower:
                    cls.last_error_type = "RATE_LIMIT"
                else:
                    cls.last_error_type = "QUOTA_EXCEEDED"
            elif "quota" in msg_lower:
                cls.last_error_type = "QUOTA_EXCEEDED"
            else:
                # Default fallback if unknown
                cls.last_error_type = "QUOTA_EXCEEDED"

            # [FIX] If model_id is provided, immediately mark it as exhausted so rotation logic knows
            if model_id and cls.last_error_type == "QUOTA_EXCEEDED":
                cls.current_profile_exhausted_models.add(model_id.lower())
                logger.warning(f"⛔ Identified specific model exhaustion: {model_id}")

            logger.critical(
                f"⛔ GLOBAL ALERT: Quota Exceeded! Type: {cls.last_error_type} (Event Signal Sent)"
            )

    @classmethod
    def reset_quota_status(cls):
        """
        Resets the global quota exceeded flag.
        """
        cls.IS_QUOTA_EXCEEDED = False
        cls.NEEDS_ROTATION = False  # Reset soft flag too
        cls.QUOTA_EXCEEDED_TIMESTAMP = 0.0
        cls.last_error_type = None
        cls.QUOTA_EXCEEDED_EVENT.clear()

        # [QUOTA-02] Reset model usage stats
        cls.current_profile_model_usage.clear()
        cls.current_profile_exhausted_models.clear()

        logger.info("✅ GLOBAL ALERT: Quota status manually reset.")

    @classmethod
    def increment_token_count(cls, count: int, model_id: str = "default"):
        """
        Increments the token count for the current profile and checks if it exceeds the limit.
        [GR-01] Implements Graceful Rotation logic.
        """
        if count <= 0:
            return

        # Ensure model_id is a valid string for key usage
        safe_model_id = model_id if model_id else "default"
        model_key = safe_model_id.lower()

        cls.current_profile_model_usage[model_key] += count
        current_usage = cls.current_profile_model_usage[model_key]

        # Retrieve limit (fallback to global hard limit)
        limit = MODEL_QUOTA_LIMITS.get(model_key, QUOTA_HARD_LIMIT)

        # Check Hard Limit (Emergency Kill / Model Exhaustion)
        if current_usage >= limit:
            logger.critical(
                f"⛔ HARD LIMIT REACHED ({model_key}): {current_usage} >= {limit}. Marking model as exhausted."
            )
            cls.current_profile_exhausted_models.add(model_key)
            # Trigger global rotation signal
            cls.set_quota_exceeded(message=f"Quota exceeded for model {model_key}")
            # Raise exception to propagate up to request processor
            raise QuotaExceededError(
                f"Quota exceeded for model {model_key} ({current_usage} >= {limit})"
            )

        # Check Soft Limit (Graceful Signal)
        # Note: Using global soft limit as baseline for rotation signal
        if current_usage >= QUOTA_SOFT_LIMIT and not cls.NEEDS_ROTATION:
            logger.warning(
                f"🔄 SOFT LIMIT REACHED ({model_key}): {current_usage} >= {QUOTA_SOFT_LIMIT}. Setting NEEDS_ROTATION flag."
            )
            cls.NEEDS_ROTATION = True

        # Log status
        limit_str = f"{QUOTA_SOFT_LIMIT}(Soft)/{limit}(Hard)"
        logger.info(
            f"📊 Token usage updated ({model_key}): +{count} => {current_usage} (Limits: {limit_str}) | Rotation Pending: {cls.NEEDS_ROTATION}"
        )