proxy_manager.py

"""
Proxy Pool Manager (Feature 13)
===============================
Multi-provider proxy management with rotation, health checks, and failover.

Synergizes with:
- Feature 4 (Scheduled): Long-running jobs rotate through pool
- Feature 15 (Smart Cache): Cache hits reduce proxy bandwidth usage
"""

import logging
import time
import random
import requests
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from enum import Enum
from threading import Lock

logger = logging.getLogger(__name__)


class RotationStrategy(Enum):
    """Proxy rotation strategies."""
    ROUND_ROBIN = "round_robin"
    LEAST_USED = "least_used"
    RANDOM = "random"
    FAILOVER = "failover"  # Use primary until it fails


class ProxyHealth(Enum):
    """Proxy health status."""
    HEALTHY = "healthy"
    DEGRADED = "degraded"  # Slow but working
    UNHEALTHY = "unhealthy"
    UNKNOWN = "unknown"


@dataclass
class ProxyProvider:
    """Configuration for a single proxy provider."""
    name: str
    url: str
    bandwidth_limit_mb: int = 1000
    priority: int = 1  # Lower = higher priority for failover
    enabled: bool = True
    use_for: List[str] = field(default_factory=list)  # Empty = use for all

    # Runtime stats (not from config)
    bandwidth_used_mb: float = 0.0
    request_count: int = 0
    error_count: int = 0
    last_error: Optional[str] = None
    last_used: Optional[datetime] = None
    last_health_check: Optional[datetime] = None
    health: ProxyHealth = ProxyHealth.UNKNOWN
    avg_response_time_ms: float = 0.0

    @property
    def is_available(self) -> bool:
        """Check if proxy is available for use."""
        if not self.enabled:
            return False
        if self.health == ProxyHealth.UNHEALTHY:
            return False
        if self.bandwidth_used_mb >= self.bandwidth_limit_mb:
            return False
        return True

    @property
    def bandwidth_remaining_mb(self) -> float:
        return max(0, self.bandwidth_limit_mb - self.bandwidth_used_mb)

    @property
    def bandwidth_percent_used(self) -> float:
        if self.bandwidth_limit_mb <= 0:
            return 0
        return (self.bandwidth_used_mb / self.bandwidth_limit_mb) * 100

    @property
    def error_rate(self) -> float:
        if self.request_count == 0:
            return 0
        return (self.error_count / self.request_count) * 100


@dataclass
class ProxyPoolConfig:
    """Configuration for the proxy pool."""
    enabled: bool = True
    rotation_strategy: RotationStrategy = RotationStrategy.ROUND_ROBIN
    health_check_interval_sec: int = 300
    health_check_url: str = "https://httpbin.org/ip"
    health_check_timeout_sec: int = 10
    max_retries: int = 3
    retry_delay_sec: float = 1.0
    auto_disable_on_limit: bool = True
    providers: List[ProxyProvider] = field(default_factory=list)


class ProxyPoolManager:
    """
    Manages a pool of proxy providers with intelligent rotation.

    Features:
    - Multiple proxy providers with priority ordering
    - Automatic rotation strategies
    - Health checks and automatic failover
    - Per-provider bandwidth tracking
    - Geographic proxy selection (optional)
    """

    def __init__(self, config: Optional[ProxyPoolConfig] = None):
        """
        Initialize proxy pool manager.

        Args:
            config: Optional ProxyPoolConfig, creates empty pool if not provided
        """
        self.config = config or ProxyPoolConfig()
        self._current_index = 0
        self._lock = Lock()

        # Global stats
        self.stats = {
            'total_requests': 0,
            'total_errors': 0,
            'total_bandwidth_mb': 0.0,
            'failovers': 0,
            'health_checks': 0,
        }

    def add_provider(self, provider: ProxyProvider):
        """Add a proxy provider to the pool."""
        self.config.providers.append(provider)
        logger.info(f"Added proxy provider: {provider.name}")

    def remove_provider(self, name: str) -> bool:
        """Remove a proxy provider by name."""
        for i, p in enumerate(self.config.providers):
            if p.name == name:
                del self.config.providers[i]
                logger.info(f"Removed proxy provider: {name}")
                return True
        return False

    def get_provider(self, name: str) -> Optional[ProxyProvider]:
        """Get a provider by name."""
        for p in self.config.providers:
            if p.name == name:
                return p
        return None

    def _get_available_providers(self, query_type: Optional[str] = None) -> List[ProxyProvider]:
        """Get list of available providers, optionally filtered by query type."""
        available = []
        for p in self.config.providers:
            if not p.is_available:
                continue
            # If provider has use_for list, check if query_type matches
            if p.use_for and query_type:
                if not any(ut.lower() in query_type.lower() for ut in p.use_for):
                    continue
            available.append(p)
        return available

    def get_proxy(self, query_type: Optional[str] = None) -> Optional[Dict[str, str]]:
        """
        Get the next proxy to use based on rotation strategy.

        Args:
            query_type: Optional query type for provider filtering

        Returns:
            Proxy dict {"server": "http://..."} or None if no proxy available
        """
        if not self.config.enabled:
            return None

        with self._lock:
            available = self._get_available_providers(query_type)

            if not available:
                logger.warning("No available proxies in pool")
                return None

            provider = self._select_provider(available)

            if provider:
                provider.last_used = datetime.now()
                return {"server": provider.url}

        return None

    def _select_provider(self, available: List[ProxyProvider]) -> Optional[ProxyProvider]:
        """Select a provider based on rotation strategy."""
        if not available:
            return None

        strategy = self.config.rotation_strategy

        if strategy == RotationStrategy.ROUND_ROBIN:
            # Simple round-robin through available providers
            self._current_index = (self._current_index + 1) % len(available)
            return available[self._current_index]

        elif strategy == RotationStrategy.LEAST_USED:
            # Select provider with least bandwidth used
            return min(available, key=lambda p: p.bandwidth_used_mb)

        elif strategy == RotationStrategy.RANDOM:
            return random.choice(available)

        elif strategy == RotationStrategy.FAILOVER:
            # Use highest priority (lowest number) available
            sorted_providers = sorted(available, key=lambda p: p.priority)
            return sorted_providers[0] if sorted_providers else None

        return available[0]

    def record_request(self, provider_name: str, bandwidth_kb: float,
                       success: bool, response_time_ms: float = 0):
        """
        Record a request made through a proxy.

        Args:
            provider_name: Name of the provider used
            bandwidth_kb: Bandwidth used in KB
            success: Whether request succeeded
            response_time_ms: Response time in milliseconds
        """
        provider = self.get_provider(provider_name)
        if not provider:
            return

        with self._lock:
            provider.request_count += 1
            provider.bandwidth_used_mb += bandwidth_kb / 1024

            self.stats['total_requests'] += 1
            self.stats['total_bandwidth_mb'] += bandwidth_kb / 1024

            if success:
                # Update rolling average response time
                if provider.avg_response_time_ms == 0:
                    provider.avg_response_time_ms = response_time_ms
                else:
                    # Exponential moving average
                    provider.avg_response_time_ms = (
                        provider.avg_response_time_ms * 0.9 + response_time_ms * 0.1
                    )
            else:
                provider.error_count += 1
                self.stats['total_errors'] += 1

            # Check if bandwidth limit reached
            if (self.config.auto_disable_on_limit and
                provider.bandwidth_used_mb >= provider.bandwidth_limit_mb):
                logger.warning(
                    f"Proxy {provider_name} reached bandwidth limit "
                    f"({provider.bandwidth_limit_mb}MB), disabling"
                )
                provider.enabled = False

    def record_error(self, provider_name: str, error: str):
        """Record an error for a provider."""
        provider = self.get_provider(provider_name)
        if provider:
            with self._lock:
                provider.error_count += 1
                provider.last_error = error
                self.stats['total_errors'] += 1

                # Mark unhealthy if too many errors
                if provider.error_rate > 50 and provider.request_count >= 10:
                    provider.health = ProxyHealth.UNHEALTHY
                    logger.warning(f"Proxy {provider_name} marked unhealthy: {error}")
                    self.stats['failovers'] += 1

    def health_check(self, provider_name: Optional[str] = None) -> Dict[str, ProxyHealth]:
        """
        Perform health check on one or all providers.

        Args:
            provider_name: Optional specific provider to check

        Returns:
            Dict mapping provider names to health status
        """
        results = {}
        providers = [self.get_provider(provider_name)] if provider_name else self.config.providers

        for provider in providers:
            if not provider or not provider.enabled:
                continue

            try:
                start = time.time()
                response = requests.get(
                    self.config.health_check_url,
                    proxies={"http": provider.url, "https": provider.url},
                    timeout=self.config.health_check_timeout_sec
                )
                elapsed_ms = (time.time() - start) * 1000

                self.stats['health_checks'] += 1
                provider.last_health_check = datetime.now()

                if response.status_code == 200:
                    if elapsed_ms < 3000:
                        provider.health = ProxyHealth.HEALTHY
                    else:
                        provider.health = ProxyHealth.DEGRADED
                else:
                    provider.health = ProxyHealth.UNHEALTHY

                results[provider.name] = provider.health
                logger.debug(f"Health check {provider.name}: {provider.health.value} ({elapsed_ms:.0f}ms)")

            except Exception as e:
                provider.health = ProxyHealth.UNHEALTHY
                provider.last_error = str(e)
                results[provider.name] = ProxyHealth.UNHEALTHY
                logger.warning(f"Health check failed for {provider.name}: {e}")

        return results

    def check_all_health(self):
        """Run health check on all providers."""
        return self.health_check()

    def get_best_provider(self, query_type: Optional[str] = None) -> Optional[ProxyProvider]:
        """
        Get the best available provider based on health and bandwidth.

        Args:
            query_type: Optional query type for filtering

        Returns:
            Best provider or None
        """
        available = self._get_available_providers(query_type)

        if not available:
            return None

        # Score providers: health, bandwidth remaining, response time
        def score(p: ProxyProvider) -> float:
            health_score = {
                ProxyHealth.HEALTHY: 100,
                ProxyHealth.DEGRADED: 50,
                ProxyHealth.UNKNOWN: 25,
                ProxyHealth.UNHEALTHY: 0
            }.get(p.health, 0)

            bandwidth_score = (p.bandwidth_remaining_mb / p.bandwidth_limit_mb) * 100
            speed_score = max(0, 100 - (p.avg_response_time_ms / 50))  # Penalize slow proxies

            return health_score * 0.4 + bandwidth_score * 0.4 + speed_score * 0.2

        return max(available, key=score)

    def reset_provider(self, name: str):
        """Reset a provider's stats and re-enable it."""
        provider = self.get_provider(name)
        if provider:
            provider.bandwidth_used_mb = 0.0
            provider.request_count = 0
            provider.error_count = 0
            provider.last_error = None
            provider.health = ProxyHealth.UNKNOWN
            provider.enabled = True
            logger.info(f"Reset proxy provider: {name}")

    def reset_all(self):
        """Reset all providers' stats."""
        for provider in self.config.providers:
            self.reset_provider(provider.name)

    def get_pool_stats(self) -> Dict[str, Any]:
        """Get comprehensive pool statistics."""
        providers_info = []
        for p in self.config.providers:
            providers_info.append({
                'name': p.name,
                'enabled': p.enabled,
                'health': p.health.value,
                'bandwidth_used_mb': round(p.bandwidth_used_mb, 2),
                'bandwidth_limit_mb': p.bandwidth_limit_mb,
                'bandwidth_percent': round(p.bandwidth_percent_used, 1),
                'request_count': p.request_count,
                'error_count': p.error_count,
                'error_rate': round(p.error_rate, 1),
                'avg_response_ms': round(p.avg_response_time_ms, 0),
                'priority': p.priority,
            })

        total_bandwidth_limit = sum(p.bandwidth_limit_mb for p in self.config.providers)
        total_bandwidth_used = sum(p.bandwidth_used_mb for p in self.config.providers)

        return {
            'enabled': self.config.enabled,
            'strategy': self.config.rotation_strategy.value,
            'provider_count': len(self.config.providers),
            'available_count': len(self._get_available_providers()),
            'total_bandwidth_limit_mb': total_bandwidth_limit,
            'total_bandwidth_used_mb': round(total_bandwidth_used, 2),
            'total_bandwidth_percent': round(
                (total_bandwidth_used / total_bandwidth_limit * 100)
                if total_bandwidth_limit > 0 else 0, 1
            ),
            'total_requests': self.stats['total_requests'],
            'total_errors': self.stats['total_errors'],
            'failovers': self.stats['failovers'],
            'health_checks': self.stats['health_checks'],
            'providers': providers_info,
        }

    def should_run_health_check(self) -> bool:
        """Check if health check should be run based on interval."""
        for provider in self.config.providers:
            if not provider.last_health_check:
                return True
            elapsed = (datetime.now() - provider.last_health_check).total_seconds()
            if elapsed >= self.config.health_check_interval_sec:
                return True
        return False

    def health_check_provider(self, provider: ProxyProvider) -> Dict[str, Any]:
        """
        Perform health check on a single provider.

        Args:
            provider: ProxyProvider to check

        Returns:
            Dict with success, response_time_ms, detected_ip, error
        """
        result = {
            'success': False,
            'response_time_ms': 0,
            'detected_ip': None,
            'error': None,
        }

        if not provider.url:
            result['error'] = 'No proxy URL configured'
            return result

        try:
            start = time.time()
            response = requests.get(
                self.config.health_check_url,
                proxies={"http": provider.url, "https": provider.url},
                timeout=self.config.health_check_timeout_sec
            )
            elapsed_ms = (time.time() - start) * 1000

            self.stats['health_checks'] += 1
            provider.last_health_check = datetime.now()
            result['response_time_ms'] = elapsed_ms

            if response.status_code == 200:
                result['success'] = True
                # Try to extract IP from response
                try:
                    data = response.json()
                    result['detected_ip'] = data.get('origin', data.get('ip', 'Unknown'))
                except:
                    result['detected_ip'] = 'Unknown'

                if elapsed_ms < 3000:
                    provider.health = ProxyHealth.HEALTHY
                else:
                    provider.health = ProxyHealth.DEGRADED
            else:
                provider.health = ProxyHealth.UNHEALTHY
                result['error'] = f"HTTP {response.status_code}"

        except requests.exceptions.Timeout:
            provider.health = ProxyHealth.UNHEALTHY
            result['error'] = 'Connection timeout'
        except requests.exceptions.ConnectionError as e:
            provider.health = ProxyHealth.UNHEALTHY
            result['error'] = 'Connection failed'
        except Exception as e:
            provider.health = ProxyHealth.UNHEALTHY
            result['error'] = str(e)[:50]

        return result

    def get_pool_status(self) -> Dict[str, Any]:
        """
        Get pool status for CLI display.

        Returns:
            Dict with active_count, total_bandwidth_mb, providers list
        """
        providers_info = []
        for p in self.config.providers:
            providers_info.append({
                'name': p.name,
                'is_available': p.is_available,
                'health': p.health.value,
                'bandwidth_used_mb': round(p.bandwidth_used_mb, 2),
                'bandwidth_limit_mb': p.bandwidth_limit_mb,
                'request_count': p.request_count,
                'error_count': p.error_count,
            })

        total_bandwidth = sum(p.bandwidth_used_mb for p in self.config.providers)
        active_count = len([p for p in self.config.providers if p.is_available])

        return {
            'active_count': active_count,
            'total_bandwidth_mb': round(total_bandwidth, 2),
            'providers': providers_info,
        }


def create_proxy_pool(config_dict: Optional[Dict] = None) -> ProxyPoolManager:
    """
    Factory function to create ProxyPoolManager from config dict.

    Args:
        config_dict: Configuration from YAML file

    Returns:
        Configured ProxyPoolManager
    """
    if not config_dict:
        return ProxyPoolManager()

    # Parse rotation strategy
    strategy_str = config_dict.get('rotation_strategy', 'round_robin')
    try:
        strategy = RotationStrategy(strategy_str)
    except ValueError:
        strategy = RotationStrategy.ROUND_ROBIN

    # Parse providers
    providers = []
    for p_dict in config_dict.get('providers', []):
        provider = ProxyProvider(
            name=p_dict.get('name', f"proxy_{len(providers)}"),
            url=p_dict.get('url', ''),
            bandwidth_limit_mb=p_dict.get('bandwidth_limit_mb', 1000),
            priority=p_dict.get('priority', 1),
            enabled=p_dict.get('enabled', True),
            use_for=p_dict.get('use_for', [])
        )
        if provider.url:  # Only add if URL is provided
            providers.append(provider)

    config = ProxyPoolConfig(
        enabled=config_dict.get('enabled', True),
        rotation_strategy=strategy,
        health_check_interval_sec=config_dict.get('health_check_interval', 300),
        health_check_url=config_dict.get('health_check_url', 'https://httpbin.org/ip'),
        health_check_timeout_sec=config_dict.get('health_check_timeout', 10),
        max_retries=config_dict.get('max_retries', 3),
        retry_delay_sec=config_dict.get('retry_delay', 1.0),
        auto_disable_on_limit=config_dict.get('auto_disable_on_limit', True),
        providers=providers
    )

    return ProxyPoolManager(config)


# Example config for YAML:
EXAMPLE_CONFIG = """
proxy_pool:
  enabled: true
  rotation_strategy: round_robin  # round_robin, least_used, random, failover
  health_check_interval: 300
  health_check_timeout: 10
  max_retries: 3
  auto_disable_on_limit: true

  providers:
    - name: webshare_free
      url: "http://username:password@p.webshare.io:80"
      bandwidth_limit_mb: 1000
      priority: 1

    - name: brightdata_backup
      url: "http://username:password@brd.superproxy.io:22225"
      bandwidth_limit_mb: 5000
      priority: 2

    - name: residential_premium
      url: "http://username:password@pr.oxylabs.io:7777"
      bandwidth_limit_mb: 10000
      priority: 3
      use_for: ["high_value", "premium"]
"""