proxy_config.py

"""Proxy and rate limiting configuration for web host restrictions"""
import asyncio
import time
from collections import defaultdict
from typing import Optional

from loguru import logger


class RateLimiter:
    """Rate limiter to prevent hitting web host restrictions"""
    
    def __init__(self, max_requests: int = 10, time_window: int = 60):
        """
        Initialize rate limiter
        
        Args:
            max_requests: Maximum requests allowed in time window
            time_window: Time window in seconds
        """
        self.max_requests = max_requests
        self.time_window = time_window
        self.requests = defaultdict(list)
        self.lock = asyncio.Lock()
    
    async def wait_if_needed(self, key: str = "default"):
        """Wait if rate limit would be exceeded"""
        async with self.lock:
            now = time.time()
            
            # Remove old requests outside time window
            self.requests[key] = [
                req_time for req_time in self.requests[key]
                if now - req_time < self.time_window
            ]
            
            # Check if we need to wait
            if len(self.requests[key]) >= self.max_requests:
                oldest_request = self.requests[key][0]
                wait_time = self.time_window - (now - oldest_request)
                
                if wait_time > 0:
                    logger.info(f"Rate limit reached for {key}, waiting {wait_time:.1f}s")
                    await asyncio.sleep(wait_time)
                    
                    # Remove the oldest request after waiting
                    now = time.time()
                    self.requests[key] = [
                        req_time for req_time in self.requests[key]
                        if now - req_time < self.time_window
                    ]
            
            # Record this request
            self.requests[key].append(now)


class ProxyManager:
    """Manage proxy rotation to bypass restrictions"""
    
    def __init__(self):
        """Initialize proxy manager"""
        self.proxies = []
        self.current_index = 0
        self.failed_proxies = set()
        
        # Load proxies from config if available
        self._load_proxies()
    
    def _load_proxies(self):
        """Load proxy list from configuration"""
        # TODO: Load from environment or config file
        # For now, using direct connection
        self.proxies = []
    
    def get_next_proxy(self) -> Optional[dict]:
        """Get next available proxy"""
        if not self.proxies:
            return None
        
        # Find next working proxy
        attempts = 0
        while attempts < len(self.proxies):
            proxy = self.proxies[self.current_index]
            self.current_index = (self.current_index + 1) % len(self.proxies)
            
            if proxy not in self.failed_proxies:
                return {'http://': proxy, 'https://': proxy}
            
            attempts += 1
        
        return None
    
    def mark_proxy_failed(self, proxy: str):
        """Mark a proxy as failed"""
        self.failed_proxies.add(proxy)
        logger.warning(f"Proxy marked as failed: {proxy}")
    
    def reset_failed_proxies(self):
        """Reset failed proxies (retry after some time)"""
        self.failed_proxies.clear()


class RetryHandler:
    """Handle retries with exponential backoff"""
    
    def __init__(self, max_retries: int = 3, base_delay: float = 1.0):
        """
        Initialize retry handler
        
        Args:
            max_retries: Maximum number of retries
            base_delay: Base delay in seconds (doubles each retry)
        """
        self.max_retries = max_retries
        self.base_delay = base_delay
    
    async def execute_with_retry(self, func, *args, **kwargs):
        """Execute function with exponential backoff retry"""
        last_exception = None
        
        for attempt in range(self.max_retries):
            try:
                return await func(*args, **kwargs)
            
            except Exception as e:
                last_exception = e
                
                if attempt < self.max_retries - 1:
                    delay = self.base_delay * (2 ** attempt)
                    logger.warning(
                        f"Attempt {attempt + 1}/{self.max_retries} failed: {e}. "
                        f"Retrying in {delay}s..."
                    )
                    await asyncio.sleep(delay)
                else:
                    logger.error(f"All {self.max_retries} attempts failed")
        
        raise last_exception


# Global instances
rate_limiter = RateLimiter(max_requests=10, time_window=60)
proxy_manager = ProxyManager()
retry_handler = RetryHandler(max_retries=3)


async def make_request_with_protection(client, method: str, url: str, **kwargs):
    """
    Make HTTP request with rate limiting, proxy support, and retry logic
    
    Args:
        client: httpx.AsyncClient instance
        method: HTTP method (get, post, etc.)
        url: Target URL
        **kwargs: Additional arguments for the request
    
    Returns:
        httpx.Response object
    """
    # Apply rate limiting
    await rate_limiter.wait_if_needed(url)
    
    # Get proxy if available
    proxy = proxy_manager.get_next_proxy()
    if proxy:
        kwargs['proxies'] = proxy
    
    # Execute with retry
    async def make_request():
        func = getattr(client, method.lower())
        return await func(url, **kwargs)
    
    try:
        return await retry_handler.execute_with_retry(make_request)
    
    except Exception as e:
        logger.error(f"Request to {url} failed after all retries: {e}")
        raise


# Convenience functions
async def get_with_protection(url: str, client=None, **kwargs):
    """GET request with all protections"""
    import httpx
    
    if client is None:
        async with httpx.AsyncClient() as client:
            return await make_request_with_protection(client, 'GET', url, **kwargs)
    else:
        return await make_request_with_protection(client, 'GET', url, **kwargs)


async def post_with_protection(url: str, client=None, **kwargs):
    """POST request with all protections"""
    import httpx
    
    if client is None:
        async with httpx.AsyncClient() as client:
            return await make_request_with_protection(client, 'POST', url, **kwargs)
    else:
        return await make_request_with_protection(client, 'POST', url, **kwargs)