"""
DAARION Platform - Rate Limiter
Protects agents from spam and abuse
"""
import time
from collections import defaultdict
from typing import Dict, Tuple, Optional
import structlog

logger = structlog.get_logger()


class RateLimiter:
    """Token bucket rate limiter per user/agent"""
    
    def __init__(self, default_rpm: int = 30, default_burst: int = 5):
        self.default_rpm = default_rpm
        self.default_burst = default_burst
        # {user_id: {agent_id: (tokens, last_update)}}
        self._buckets: Dict[str, Dict[str, Tuple[float, float]]] = defaultdict(dict)
        # {agent_id: rpm_limit}
        self._agent_limits: Dict[str, int] = {}
        # Blocked users (temporary)
        self._blocked: Dict[str, float] = {}  # user_id: unblock_time
    
    def set_agent_limit(self, agent_id: str, rpm: int):
        """Set rate limit for specific agent"""
        self._agent_limits[agent_id] = rpm
        logger.info("rate_limit_set", agent_id=agent_id, rpm=rpm)
    
    def _get_limit(self, agent_id: str) -> int:
        """Get RPM limit for agent"""
        return self._agent_limits.get(agent_id, self.default_rpm)
    
    def _refill_tokens(self, user_id: str, agent_id: str) -> float:
        """Refill tokens based on time elapsed"""
        now = time.time()
        
        if agent_id not in self._buckets[user_id]:
            # New user - full bucket
            self._buckets[user_id][agent_id] = (self.default_burst, now)
            return self.default_burst
        
        tokens, last_update = self._buckets[user_id][agent_id]
        rpm = self._get_limit(agent_id)
        
        # Calculate tokens to add (tokens per second = rpm / 60)
        elapsed = now - last_update
        tokens_to_add = elapsed * (rpm / 60.0)
        
        # Cap at burst limit
        new_tokens = min(self.default_burst, tokens + tokens_to_add)
        self._buckets[user_id][agent_id] = (new_tokens, now)
        
        return new_tokens
    
    def is_blocked(self, user_id: str) -> bool:
        """Check if user is temporarily blocked"""
        if user_id not in self._blocked:
            return False
        if time.time() > self._blocked[user_id]:
            del self._blocked[user_id]
            return False
        return True
    
    def block_user(self, user_id: str, seconds: int = 60):
        """Temporarily block user"""
        self._blocked[user_id] = time.time() + seconds
        logger.warning("user_blocked", user_id=user_id, seconds=seconds)
    
    def check(self, user_id: str, agent_id: str) -> Tuple[bool, Optional[str]]:
        """
        Check if request is allowed.
        
        Returns:
            (allowed, error_message)
        """
        # Check block list
        if self.is_blocked(user_id):
            remaining = int(self._blocked[user_id] - time.time())
            return False, f"Занадто багато запитів. Спробуйте через {remaining} сек."
        
        # Refill and check tokens
        tokens = self._refill_tokens(user_id, agent_id)
        
        if tokens < 1:
            # No tokens - calculate wait time
            rpm = self._get_limit(agent_id)
            wait_seconds = int(60 / rpm)
            logger.warning("rate_limit_exceeded", user_id=user_id, agent_id=agent_id)
            return False, f"Ліміт запитів. Спробуйте через {wait_seconds} сек."
        
        # Consume token
        self._buckets[user_id][agent_id] = (tokens - 1, time.time())
        return True, None
    
    def get_remaining(self, user_id: str, agent_id: str) -> int:
        """Get remaining requests for user"""
        tokens = self._refill_tokens(user_id, agent_id)
        return int(tokens)


# Singleton instance
_limiter = None

def get_rate_limiter() -> RateLimiter:
    """Get or create rate limiter instance"""
    global _limiter
    if _limiter is None:
        _limiter = RateLimiter()
        # Set agent-specific limits
        _limiter.set_agent_limit("helion", 60)
        _limiter.set_agent_limit("nutra", 30)
        _limiter.set_agent_limit("greenfood", 20)
        _limiter.set_agent_limit("druid", 20)
        _limiter.set_agent_limit("daarwizz", 100)
    return _limiter
