"""
Rate Limiting and Throttling for Connectors
Token bucket and sliding window rate limiting algorithms
"""

import time
import threading
from typing import Dict, Any, Optional, Callable
from dataclasses import dataclass
from collections import deque
from functools import wraps
from .exceptions import RateLimitError
from .logging_config import get_logger, log_performance

logger = get_logger(__name__)


@dataclass
class RateLimitConfig:
    """Configuration for rate limiting"""
    
    requests_per_period: int
    period_seconds: float
    
    # Token bucket specific
    burst_capacity: Optional[int] = None
    refill_rate: Optional[float] = None
    
    # Sliding window specific
    window_size: Optional[float] = None
    
    # Behavior when limit is exceeded
    block_when_exceeded: bool = True
    retry_after_seconds: Optional[float] = None
    
    def __post_init__(self):
        """Set default values"""
        if self.burst_capacity is None:
            self.burst_capacity = self.requests_per_period
        
        if self.refill_rate is None:
            self.refill_rate = self.requests_per_period / self.period_seconds
        
        if self.window_size is None:
            self.window_size = self.period_seconds
        
        if self.retry_after_seconds is None:
            self.retry_after_seconds = self.period_seconds


class TokenBucket:
    """Token bucket rate limiter"""
    
    def __init__(self, config: RateLimitConfig):
        self.config = config
        self.tokens = config.burst_capacity
        self.last_refill = time.time()
        self._lock = threading.Lock()
    
    def _refill_tokens(self) -> None:
        """Refill tokens based on elapsed time"""
        current_time = time.time()
        elapsed = current_time - self.last_refill
        
        if elapsed > 0:
            tokens_to_add = elapsed * self.config.refill_rate
            self.tokens = min(self.config.burst_capacity, self.tokens + tokens_to_add)
            self.last_refill = current_time
    
    def consume(self, tokens: int = 1) -> bool:
        """Consume tokens if available"""
        with self._lock:
            self._refill_tokens()
            
            if self.tokens >= tokens:
                self.tokens -= tokens
                return True
            
            return False
    
    def wait_for_token(self, tokens: int = 1, timeout: Optional[float] = None) -> bool:
        """Wait for tokens to become available"""
        start_time = time.time()
        
        while True:
            if self.consume(tokens):
                return True
            
            if timeout and (time.time() - start_time) >= timeout:
                return False
            
            # Calculate wait time
            wait_time = tokens / self.config.refill_rate
            time.sleep(min(wait_time, 0.1))  # Sleep in small increments
    
    def get_available_tokens(self) -> int:
        """Get number of available tokens"""
        with self._lock:
            self._refill_tokens()
            return int(self.tokens)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get rate limiter statistics"""
        with self._lock:
            self._refill_tokens()
            return {
                'available_tokens': int(self.tokens),
                'burst_capacity': self.config.burst_capacity,
                'refill_rate': self.config.refill_rate,
                'last_refill': self.last_refill
            }


class SlidingWindow:
    """Sliding window rate limiter"""
    
    def __init__(self, config: RateLimitConfig):
        self.config = config
        self.requests = deque()
        self._lock = threading.Lock()
    
    def _cleanup_old_requests(self) -> None:
        """Remove requests outside the window"""
        current_time = time.time()
        cutoff_time = current_time - self.config.window_size
        
        while self.requests and self.requests[0] <= cutoff_time:
            self.requests.popleft()
    
    def is_allowed(self) -> bool:
        """Check if request is allowed"""
        with self._lock:
            self._cleanup_old_requests()
            
            return len(self.requests) < self.config.requests_per_period
    
    def record_request(self) -> None:
        """Record a request timestamp"""
        with self._lock:
            self.requests.append(time.time())
    
    def get_wait_time(self) -> float:
        """Get time to wait before next request is allowed"""
        with self._lock:
            self._cleanup_old_requests()
            
            if len(self.requests) < self.config.requests_per_period:
                return 0.0
            
            # Calculate when oldest request will fall out of window
            oldest_request = self.requests[0]
            wait_time = (oldest_request + self.config.window_size) - time.time()
            
            return max(0.0, wait_time)
    
    def get_stats(self) -> Dict[str, Any]:
        """Get rate limiter statistics"""
        with self._lock:
            self._cleanup_old_requests()
            return {
                'current_requests': len(self.requests),
                'max_requests': self.config.requests_per_period,
                'window_size': self.config.window_size,
                'oldest_request': self.requests[0] if self.requests else None,
                'newest_request': self.requests[-1] if self.requests else None
            }


class RateLimiter:
    """Generic rate limiter with multiple algorithms"""
    
    def __init__(
        self,
        config: RateLimitConfig,
        algorithm: str = "token_bucket"
    ):
        self.config = config
        self.algorithm = algorithm
        
        if algorithm == "token_bucket":
            self.limiter = TokenBucket(config)
        elif algorithm == "sliding_window":
            self.limiter = SlidingWindow(config)
        else:
            raise ValueError(f"Unknown rate limiting algorithm: {algorithm}")
    
    def acquire(self, tokens: int = 1) -> bool:
        """Acquire permission to make requests"""
        if self.algorithm == "token_bucket":
            return self.limiter.consume(tokens)
        else:  # sliding_window
            if self.limiter.is_allowed():
                self.limiter.record_request()
                return True
            return False
    
    def acquire_or_wait(self, tokens: int = 1, timeout: Optional[float] = None) -> bool:
        """Acquire permission, waiting if necessary"""
        if self.algorithm == "token_bucket":
            return self.limiter.wait_for_token(tokens, timeout)
        else:  # sliding_window
            start_time = time.time()
            
            while True:
                if self.limiter.is_allowed():
                    self.limiter.record_request()
                    return True
                
                if timeout and (time.time() - start_time) >= timeout:
                    return False
                
                wait_time = self.limiter.get_wait_time()
                if wait_time > 0:
                    time.sleep(min(wait_time, 0.1))
    
    def check_limit(self) -> Dict[str, Any]:
        """Check current rate limit status"""
        stats = self.limiter.get_stats()
        
        if self.algorithm == "token_bucket":
            available_tokens = stats['available_tokens']
            is_allowed = available_tokens > 0
            wait_time = 0.0 if is_allowed else (1.0 / self.config.refill_rate)
        else:  # sliding_window
            current_requests = stats['current_requests']
            is_allowed = current_requests < self.config.requests_per_period
            wait_time = self.limiter.get_wait_time()
        
        return {
            'is_allowed': is_allowed,
            'wait_time': wait_time,
            'stats': stats,
            'algorithm': self.algorithm
        }
    
    def enforce_limit(self, tokens: int = 1) -> None:
        """Enforce rate limit, raising exception if exceeded"""
        if not self.acquire(tokens):
            status = self.check_limit()
            
            if self.config.block_when_exceeded:
                raise RateLimitError(
                    f"Rate limit exceeded. {self.config.requests_per_period} requests per "
                    f"{self.config.period_seconds}s allowed. Wait {status['wait_time']:.2f}s.",
                    retry_after=int(status['wait_time']) + 1
                )
            else:
                logger.warning(
                    f"Rate limit exceeded. Consider waiting {status['wait_time']:.2f}s"
                )


class MultiRateLimiter:
    """Rate limiter with multiple limits (e.g., per second, per minute, per hour)"""
    
    def __init__(self, limits: Dict[str, RateLimitConfig]):
        self.limiters = {}
        
        for name, config in limits.items():
            self.limiters[name] = RateLimiter(config)
    
    def acquire(self, tokens: int = 1) -> bool:
        """Acquire permission from all limiters"""
        for limiter in self.limiters.values():
            if not limiter.acquire(tokens):
                return False
        return True
    
    def acquire_or_wait(self, tokens: int = 1, timeout: Optional[float] = None) -> bool:
        """Acquire permission, waiting for the slowest limiter"""
        max_wait_time = 0.0
        
        for limiter in self.limiters.values():
            status = limiter.check_limit()
            if not status['is_allowed']:
                max_wait_time = max(max_wait_time, status['wait_time'])
        
        if max_wait_time > 0:
            if timeout and max_wait_time > timeout:
                return False
            time.sleep(max_wait_time)
        
        # Try again after waiting
        return self.acquire(tokens)
    
    def enforce_limit(self, tokens: int = 1) -> None:
        """Enforce all rate limits"""
        for name, limiter in self.limiters.items():
            try:
                limiter.enforce_limit(tokens)
            except RateLimitError as e:
                # Add context about which limit was exceeded
                raise RateLimitError(
                    f"Rate limit '{name}' exceeded: {e.message}",
                    retry_after=e.retry_after
                ) from e
    
    def get_status(self) -> Dict[str, Dict[str, Any]]:
        """Get status of all limiters"""
        status = {}
        for name, limiter in self.limiters.items():
            status[name] = limiter.check_limit()
        return status


def rate_limit(
    requests_per_period: int,
    period_seconds: float,
    algorithm: str = "token_bucket",
    block_when_exceeded: bool = True,
    retry_after_seconds: Optional[float] = None
):
    """Decorator for rate limiting function calls"""
    def decorator(func: Callable) -> Callable:
        # Create rate limiter
        config = RateLimitConfig(
            requests_per_period=requests_per_period,
            period_seconds=period_seconds,
            block_when_exceeded=block_when_exceeded,
            retry_after_seconds=retry_after_seconds
        )
        
        limiter = RateLimiter(config, algorithm)
        
        @wraps(func)
        @log_performance
        def wrapper(*args, **kwargs):
            # Check rate limit
            status = limiter.check_limit()
            
            if not status['is_allowed']:
                if block_when_exceeded:
                    raise RateLimitError(
                        f"Rate limit exceeded for function {func.__name__}. "
                        f"Wait {status['wait_time']:.2f}s.",
                        retry_after=int(status['wait_time']) + 1
                    )
                else:
                    logger.warning(
                        f"Rate limit exceeded for function {func.__name__}. "
                        f"Consider waiting {status['wait_time']:.2f}s"
                    )
            
            # Acquire permission
            if not limiter.acquire():
                if block_when_exceeded:
                    raise RateLimitError(
                        f"Failed to acquire rate limit for function {func.__name__}"
                    )
                else:
                    logger.warning(
                        f"Failed to acquire rate limit for function {func.__name__}"
                    )
            
            return func(*args, **kwargs)
        
        return wrapper
    return decorator


def async_rate_limit(
    requests_per_period: int,
    period_seconds: float,
    algorithm: str = "token_bucket",
    block_when_exceeded: bool = True,
    retry_after_seconds: Optional[float] = None
):
    """Async decorator for rate limiting function calls"""
    def decorator(func: Callable) -> Callable:
        # Create rate limiter
        config = RateLimitConfig(
            requests_per_period=requests_per_period,
            period_seconds=period_seconds,
            block_when_exceeded=block_when_exceeded,
            retry_after_seconds=retry_after_seconds
        )
        
        limiter = RateLimiter(config, algorithm)
        
        @wraps(func)
        async def wrapper(*args, **kwargs):
            # Check rate limit
            status = limiter.check_limit()
            
            if not status['is_allowed']:
                if block_when_exceeded:
                    raise RateLimitError(
                        f"Rate limit exceeded for async function {func.__name__}. "
                        f"Wait {status['wait_time']:.2f}s.",
                        retry_after=int(status['wait_time']) + 1
                    )
                else:
                    logger.warning(
                        f"Rate limit exceeded for async function {func.__name__}. "
                        f"Consider waiting {status['wait_time']:.2f}s"
                    )
            
            # Acquire permission (with waiting)
            if not limiter.acquire_or_wait():
                if block_when_exceeded:
                    raise RateLimitError(
                        f"Failed to acquire rate limit for async function {func.__name__}"
                    )
                else:
                    logger.warning(
                        f"Failed to acquire rate limit for async function {func.__name__}"
                    )
            
            return await func(*args, **kwargs)
        
        return wrapper
    return decorator


# Global rate limiter registry
_global_rate_limiters: Dict[str, RateLimiter] = {}


def register_rate_limiter(name: str, config: RateLimitConfig, algorithm: str = "token_bucket") -> RateLimiter:
    """Register a global rate limiter"""
    limiter = RateLimiter(config, algorithm)
    _global_rate_limiters[name] = limiter
    logger.info(f"Registered rate limiter '{name}' with {algorithm} algorithm")
    return limiter


def get_rate_limiter(name: str) -> Optional[RateLimiter]:
    """Get a registered rate limiter"""
    return _global_rate_limiters.get(name)


def enforce_global_rate_limit(name: str, tokens: int = 1) -> None:
    """Enforce a global rate limit"""
    limiter = get_rate_limiter(name)
    if limiter:
        limiter.enforce_limit(tokens)
    else:
        logger.warning(f"Rate limiter '{name}' not found")
