"""
FastRDP-NG: Proxy Manager Module.
Fetches free SOCKS5/HTTP proxy lists, tests them, and provides rotating proxy
support for all outbound connections. Integrates with scanner + spray modules.
"""

import asyncio
import logging
import random
import time
from dataclasses import dataclass, field
from typing import List, Optional, Tuple, Callable

import aiohttp
from aiohttp_socks import open_connection as socks_open_connection
from aiohttp_socks import ProxyType, ProxyConnectionError

logger = logging.getLogger("FastRDP-NG")

# ── Proxy sources (free proxy lists from GitHub) ─────────────────────────
PROXY_SOURCES = [
    # Proxifly format: JSON array of {ip, port, protocol, ...}
    "https://raw.githubusercontent.com/proxifly/free-proxy-list/main/proxies.json",
    # TheSpeedX format: ip:port per line (HTTP)
    "https://raw.githubusercontent.com/TheSpeedX/PROXY-List/master/http.txt",
    # TheSpeedX format: ip:port per line (SOCKS4)
    "https://raw.githubusercontent.com/TheSpeedX/PROXY-List/master/socks4.txt",
    # TheSpeedX format: ip:port per line (SOCKS5)
    "https://raw.githubusercontent.com/TheSpeedX/PROXY-List/master/socks5.txt",
    # ProxyScrape format: ip:port per line
    "https://api.proxyscrape.com/v2/?request=displayproxies&protocol=socks5&timeout=10000&country=all",
    "https://api.proxyscrape.com/v2/?request=displayproxies&protocol=http&timeout=10000&country=all",
    # Monosans format: ip:port per line
    "https://raw.githubusercontent.com/monosans/proxy-list/main/proxies/all.txt",
]

# Test URL for proxy verification (fast, reliable endpoint)
PROXY_TEST_URL = "http://httpbin.org/ip"
PROXY_TEST_TIMEOUT = 8.0

# Default credentials for proxy testing (most free proxies don't use auth)
NO_AUTH = ("", "")

# ── How often to refresh from sources (seconds) ──────────────────────────
REFRESH_INTERVAL = 300  # 5 minutes


@dataclass
class ProxyEntry:
    """A single proxy entry with status tracking."""
    host: str
    port: int
    protocol: str          # "socks5", "socks4", "http"
    alive: bool = False
    latency: float = 0.0   # ms
    last_tested: float = 0.0
    failures: int = 0
    username: str = ""
    password: str = ""

    @property
    def proxy_type(self) -> ProxyType:
        if self.protocol == "socks4":
            return ProxyType.SOCKS4
        elif self.protocol == "socks5":
            return ProxyType.SOCKS5
        else:
            return ProxyType.HTTP

    @property
    def url(self) -> str:
        """Return proxy URL string for aiohttp_socks."""
        if self.username:
            return f"{self.protocol}://{self.username}:{self.password}@{self.host}:{self.port}"
        return f"{self.protocol}://{self.host}:{self.port}"

    def __str__(self):
        status = "\u2713" if self.alive else "\u2717"
        return f"{status} {self.protocol}://{self.host}:{self.port} ({self.latency:.0f}ms)"


class ProxyManager:
    """
    Manages a dynamic list of proxies with fetching, testing, rotation.
    Thread-safe for use across asyncio event loops.
    """

    def __init__(self):
        self.proxies: List[ProxyEntry] = []
        self._working: List[ProxyEntry] = []
        self._rotation_index = 0
        self._lock = asyncio.Lock()

        # State
        self.enabled = False
        self.auto_rotate = True
        self.testing = False
        self.fetching = False
        self.last_refresh = 0.0

        # Stats
        self.total_fetched = 0
        self.working_count = 0
        self.failed_count = 0

        # Session for HTTP requests
        self._session: Optional[aiohttp.ClientSession] = None

    async def _get_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(
                timeout=aiohttp.ClientTimeout(total=15)
            )
        return self._session

    async def close(self):
        if self._session and not self._session.closed:
            await self._session.close()

    # ── FETCHING ──────────────────────────────────────────────────────────

    async def fetch_from_sources(
        self,
        progress_callback: Optional[Callable[[str], None]] = None
    ) -> int:
        """
        Fetch proxy lists from all configured sources.
        Returns total number of unique proxies collected.
        """
        async with self._lock:
            if self.fetching:
                return len(self.proxies)
            self.fetching = True

        try:
            session = await self._get_session()
            all_proxies: List[Tuple[str, int, str]] = []  # (host, port, protocol)
            seen = set()

            for source_url in PROXY_SOURCES:
                try:
                    if progress_callback:
                        progress_callback(f"Fetching: {source_url.split('/')[-1][:40]}...")

                    async with session.get(source_url, timeout=10) as resp:
                        if resp.status != 200:
                            continue
                        text = await resp.text()

                    # Parse based on source URL pattern
                    if "proxifly" in source_url.lower() or source_url.endswith(".json"):
                        # JSON format: [{"ip":"...","port":...,"protocol":"..."}]
                        parsed = self._parse_proxifly_json(text)
                    else:
                        # Plain text: ip:port per line
                        parsed = self._parse_plain_text(text)

                    for host, port, protocol in parsed:
                        key = (host, port, protocol)
                        if key not in seen:
                            seen.add(key)
                            all_proxies.append((host, port, protocol))

                    if progress_callback:
                        progress_callback(f"  -> Got {len(parsed)} proxies from {source_url.split('/')[-1][:30]}")

                except (asyncio.TimeoutError, aiohttp.ClientError, Exception) as e:
                    logger.debug(f"Proxy source failed: {source_url[:60]} - {e}")
                    continue

            # Deduplicate and create entries
            new_proxies = []
            for host, port, protocol in all_proxies:
                # Skip duplicates with existing entries
                exists = any(
                    p.host == host and p.port == port and p.protocol == protocol
                    for p in self.proxies
                )
                if not exists:
                    new_proxies.append(ProxyEntry(
                        host=host, port=port, protocol=protocol
                    ))

            self.proxies.extend(new_proxies)
            self.total_fetched = len(self.proxies)

            if progress_callback:
                progress_callback(f"\U0001f4e1 Total: {len(self.proxies)} unique proxies collected")

            self.last_refresh = time.time()
            return len(self.proxies)

        finally:
            async with self._lock:
                self.fetching = False

    def _parse_proxifly_json(self, text: str) -> List[Tuple[str, int, str]]:
        """Parse proxifly JSON format."""
        import json
        result = []
        try:
            data = json.loads(text)
            if isinstance(data, list):
                for entry in data:
                    ip = entry.get("ip", "") or entry.get("host", "")
                    port = entry.get("port", 0)
                    protocol = entry.get("protocol", "socks5").lower()
                    if ip and port:
                        # Map protocol names
                        if protocol in ("socks5", "socks4", "http", "https"):
                            if protocol == "https":
                                protocol = "http"
                            result.append((ip, int(port), protocol))
        except (json.JSONDecodeError, ValueError):
            pass
        return result

    def _parse_plain_text(self, text: str) -> List[Tuple[str, int, str]]:
        """Parse 'ip:port' per line format."""
        result = []
        # Determine protocol from context (will be overwritten by caller)
        # We try to autodetect or default to socks5
        for line in text.splitlines():
            line = line.strip()
            if not line or line.startswith("#") or line.startswith("//"):
                continue
            # Try ip:port format
            if ":" in line and not line.startswith("["):
                parts = line.split(":")
                if len(parts) == 2:
                    ip, port_str = parts
                    try:
                        port = int(port_str)
                        if 1 <= port <= 65535:
                            # Default to socks5 (source determines actual type)
                            result.append((ip, port, "socks5"))
                    except ValueError:
                        pass
        return result

    # ── TESTING ───────────────────────────────────────────────────────────

    async def test_proxy(
        self,
        proxy: ProxyEntry,
        test_url: str = PROXY_TEST_URL,
        timeout: float = PROXY_TEST_TIMEOUT
    ) -> bool:
        """
        Test a single proxy by making an HTTP request through it.
        Returns True if proxy is working.
        """
        start = time.time()
        try:
            # Use aiohttp with ProxyConnector for testing
            from aiohttp_socks import ProxyConnector

            connector = ProxyConnector(
                proxy_type=proxy.proxy_type,
                proxy_host=proxy.host,
                proxy_port=proxy.port,
                username=proxy.username or None,
                password=proxy.password or None,
                rdns=True
            )

            async with aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=timeout)) as session:
                async with session.get(test_url, timeout=timeout) as resp:
                    if resp.status == 200:
                        proxy.alive = True
                        proxy.latency = (time.time() - start) * 1000  # ms
                        proxy.last_tested = time.time()
                        proxy.failures = 0
                        return True
                    else:
                        proxy.alive = False
                        proxy.failures += 1
                        return False

        except (ProxyConnectionError, asyncio.TimeoutError,
                aiohttp.ClientError, OSError, ConnectionError):
            proxy.alive = False
            proxy.failures += 1
            return False
        finally:
            proxy.last_tested = time.time()

    async def test_all(
        self,
        max_concurrent: int = 100,
        progress_callback: Optional[Callable[[int, int], None]] = None
    ) -> int:
        """
        Test all untested/failed proxies in the pool.
        Returns number of working proxies found.
        """
        async with self._lock:
            if self.testing:
                return self.working_count
            self.testing = True

        try:
            # Test proxies that haven't been tested or have failures
            to_test = [
                p for p in self.proxies
                if not p.alive or p.failures > 0
            ]

            if not to_test:
                # Test a random sample if all are already alive
                to_test = random.sample(self.proxies, min(50, len(self.proxies)))

            total = len(to_test)
            tested = 0
            working = 0

            sem = asyncio.Semaphore(max_concurrent)

            async def test_one(proxy: ProxyEntry):
                nonlocal tested, working
                async with sem:
                    ok = await self.test_proxy(proxy)
                    tested += 1
                    if ok:
                        working += 1
                    if progress_callback and tested % 10 == 0:
                        progress_callback(tested, total)

            tasks = [asyncio.create_task(test_one(p)) for p in to_test]
            if tasks:
                await asyncio.gather(*tasks, return_exceptions=True)

            # Update working list
            self._working = [p for p in self.proxies if p.alive]
            self.working_count = len(self._working)
            self.failed_count = len(self.proxies) - self.working_count

            return self.working_count

        finally:
            async with self._lock:
                self.testing = False

    # ── ROTATION ──────────────────────────────────────────────────────────

    async def get_proxy(self) -> Optional[ProxyEntry]:
        """
        Get the next working proxy (round-robin if auto_rotate is enabled).
        Returns None if no working proxies available.
        """
        async with self._lock:
            if not self._working:
                return None

            if self.auto_rotate and len(self._working) > 1:
                proxy = self._working[self._rotation_index % len(self._working)]
                self._rotation_index += 1
            else:
                proxy = self._working[self._rotation_index % len(self._working)]

            return proxy

    def mark_bad(self, proxy: ProxyEntry):
        """Mark a proxy as failed (connection error)."""
        proxy.failures += 1
        if proxy.failures >= 3:
            proxy.alive = False
            if proxy in self._working:
                self._working.remove(proxy)
            self.working_count = len(self._working)
            self.failed_count = len(self.proxies) - self.working_count

    # ── CONNECTION WRAPPER ────────────────────────────────────────────────

    async def open_connection(self, host: str, port: int, **kwargs):
        """
        Open a TCP connection through the current proxy.
        If proxy is disabled or no working proxies available, falls back to
        direct `asyncio.open_connection`.

        Returns (reader, writer) - same as asyncio.open_connection.
        """
        if not self.enabled:
            return await asyncio.open_connection(host, port, **kwargs)

        proxy = await self.get_proxy()
        if proxy is None:
            # No working proxy — fallback to direct
            logger.debug("No working proxy available, falling back to direct connection")
            return await asyncio.open_connection(host, port, **kwargs)

        try:
            return await asyncio.wait_for(
                socks_open_connection(
                    proxy_url=proxy.url,
                    host=host,
                    port=port,
                    **kwargs
                ),
                timeout=kwargs.get("timeout", 10.0) if "timeout" in kwargs else 10.0
            )
        except (ProxyConnectionError, asyncio.TimeoutError, OSError, ConnectionError) as e:
            self.mark_bad(proxy)
            # Retry with next proxy
            logger.debug(f"Proxy {proxy} failed: {e}, trying next...")
            return await self.open_connection(host, port, **kwargs)

    # ── STATS ─────────────────────────────────────────────────────────────

    def get_stats(self) -> dict:
        """Return proxy pool statistics."""
        return {
            "enabled": self.enabled,
            "total": len(self.proxies),
            "working": self.working_count,
            "failed": self.failed_count,
            "fetching": self.fetching,
            "testing": self.testing,
            "auto_rotate": self.auto_rotate,
            "last_refresh": self.last_refresh,
        }