import asyncio
import aiohttp
import aiohttp_socks
import random
import time
from core.log import logger


class ProxyValidator:
    """代理验证器 - 支持 HTTP/HTTPS/SOCKS4/SOCKS5"""
    
    def __init__(self, max_concurrency=50, timeout=5):
        # 验证目标源
        self.http_sources = [
            "http://httpbin.org/ip",
            "http://api.ipify.org"
        ]
        self.https_sources = [
            "https://httpbin.org/ip",
            "https://api.ipify.org"
        ]
        self.semaphore = asyncio.Semaphore(max_concurrency)
        self.timeout = timeout
        self.session = None

    async def __aenter__(self):
        """异步上下文管理器入口"""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """异步上下文管理器出口"""
        if self.session:
            await self.session.close()
            self.session = None

    def _get_test_url(self, protocol: str) -> str:
        """根据协议获取测试 URL"""
        protocol = protocol.lower()
        if protocol == 'https':
            return random.choice(self.https_sources)
        return random.choice(self.http_sources)

    def _create_connector(self, ip: str, port: int, protocol: str):
        """创建代理连接器"""
        protocol = protocol.lower()
        
        if protocol == 'socks4':
            return aiohttp_socks.ProxyConnector(
                proxy_type=aiohttp_socks.ProxyType.SOCKS4,
                host=ip,
                port=port,
                rdns=True
            )
        elif protocol == 'socks5':
            return aiohttp_socks.ProxyConnector(
                proxy_type=aiohttp_socks.ProxyType.SOCKS5,
                host=ip,
                port=port,
                rdns=True
            )
        elif protocol in ('http', 'https'):
            # HTTP/HTTPS 使用普通 connector，在请求时指定 proxy 参数
            return aiohttp.TCPConnector(ssl=False, limit=0, force_close=True)
        else:
            # 未知协议默认使用 HTTP
            return aiohttp.TCPConnector(ssl=False, limit=0, force_close=True)

    async def validate(self, ip: str, port: int, protocol: str = 'http'):
        """
        验证单个代理是否可用
        
        Args:
            ip: 代理 IP
            port: 代理端口
            protocol: 协议类型 (http/https/socks4/socks5)
            
        Returns:
            (is_valid: bool, latency_ms: float)
        """
        protocol = protocol.lower()
        test_url = self._get_test_url(protocol)
        
        async with self.semaphore:
            start_time = time.time()
            
            try:
                if protocol in ('socks4', 'socks5'):
                    return await self._validate_socks(ip, port, protocol, test_url, start_time)
                else:
                    return await self._validate_http(ip, port, protocol, test_url, start_time)
                    
            except asyncio.TimeoutError:
                logger.warning(f"验证超时: {ip}:{port} ({protocol})")
                return False, 0
            except Exception as e:
                logger.warning(f"验证失败: {ip}:{port} ({protocol}) - {e}")
                return False, 0

    async def _validate_http(self, ip: str, port: int, protocol: str, test_url: str, start_time: float):
        """验证 HTTP/HTTPS 代理"""
        proxy_url = f"http://{ip}:{port}"
        
        connector = aiohttp.TCPConnector(ssl=False, limit=0, force_close=True)
        timeout = aiohttp.ClientTimeout(total=self.timeout, connect=3)
        
        async with aiohttp.ClientSession(
            connector=connector,
            timeout=timeout
        ) as session:
            async with session.get(
                test_url,
                proxy=proxy_url,
                allow_redirects=True
            ) as response:
                if response.status in [200, 301, 302]:
                    try:
                        content = await response.text()
                        if 'ip' in content.lower() or 'origin' in content.lower():
                            latency = round((time.time() - start_time) * 1000, 2)
                            logger.info(f"验证成功: {ip}:{port} ({protocol}) - 延迟: {latency}ms")
                            return True, latency
                    except:
                        pass
                    
                    # 内容解析失败但状态码正常，也算可用
                    latency = round((time.time() - start_time) * 1000, 2)
                    logger.info(f"验证成功: {ip}:{port} ({protocol}) - 延迟: {latency}ms")
                    return True, latency
                    
                return False, 0

    async def _validate_socks(self, ip: str, port: int, protocol: str, test_url: str, start_time: float):
        """验证 SOCKS4/SOCKS5 代理"""
        proxy_type = (
            aiohttp_socks.ProxyType.SOCKS4 
            if protocol == 'socks4' 
            else aiohttp_socks.ProxyType.SOCKS5
        )
        
        connector = aiohttp_socks.ProxyConnector(
            proxy_type=proxy_type,
            host=ip,
            port=port,
            rdns=True,  # 远程 DNS 解析，避免 DNS 泄漏
            ssl=False
        )
        
        timeout = aiohttp.ClientTimeout(total=self.timeout, connect=3)
        
        try:
            async with aiohttp.ClientSession(
                connector=connector,
                timeout=timeout
            ) as session:
                async with session.get(test_url, allow_redirects=True) as response:
                    if response.status in [200, 301, 302]:
                        try:
                            content = await response.text()
                            if 'ip' in content.lower() or 'origin' in content.lower():
                                latency = round((time.time() - start_time) * 1000, 2)
                                logger.info(f"验证成功: {ip}:{port} ({protocol}) - 延迟: {latency}ms")
                                return True, latency
                        except:
                            pass
                        
                        # 内容解析失败但状态码正常
                        latency = round((time.time() - start_time) * 1000, 2)
                        logger.info(f"验证成功: {ip}:{port} ({protocol}) - 延迟: {latency}ms")
                        return True, latency
                        
                    return False, 0
        finally:
            await connector.close()


class ProxyValidatorLegacy:
    """
    兼容旧版本的验证器
    保持原有接口不变
    """
    def __init__(self, max_concurrency=50, timeout=5):
        self.validator = ProxyValidator(max_concurrency, timeout)

    async def __aenter__(self):
        await self.validator.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.validator.__aexit__(exc_type, exc_val, exc_tb)

    async def validate(self, ip, port, protocol='http'):
        return await self.validator.validate(ip, port, protocol)