"""Network protection — keep connections alive through CSCA blocking.

CSCA blocks network in several ways:
1. WFP/Firewall rules — blocks outbound connections by process name
2. DNS blocking — redirects or drops DNS queries
3. Process-based filtering — kills processes that make network requests

Our defense strategy:
- Pre-warm connections BEFORE CSCA starts blocking (established TCP survives)
- TCP keepalive on all sockets (prevents idle timeout)
- Route through localhost proxy (loopback is never blocked)
- Connection watchdog with auto-reconnect
- WFP firewall rule to whitelist our process (requires admin, optional)

Usage:
    from src.utils.network_protect import NetworkProtector

    protector = NetworkProtector(proxy_url="socks5://127.0.0.1:7890")
    protector.start()  # pre-warms connections + starts watchdog
    ...
    protector.stop()
"""

from __future__ import annotations

import ctypes
import logging
import os
import socket
import subprocess
import sys
import threading
import time
from dataclasses import dataclass, field

log = logging.getLogger(__name__)

# Endpoints to pre-warm connections to (host, port)
_WARMUP_ENDPOINTS: list[tuple[str, int]] = [
    ("api.deepseek.com", 443),
    ("dashscope.aliyuncs.com", 443),
    ("license.muravey.app", 443),
]


# ------------------------------------------------------------------ #
# TCP keepalive helper
# ------------------------------------------------------------------ #

def _set_tcp_keepalive(sock: socket.socket, idle: int = 15, interval: int = 5, count: int = 3) -> None:
    """Enable TCP keepalive on a socket.

    Args:
        idle: seconds before first keepalive probe (default 15)
        interval: seconds between probes (default 5)
        count: max failed probes before connection is dead (default 3)
    """
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)

    if sys.platform == "darwin":
        # macOS: TCP_KEEPALIVE = idle time
        TCP_KEEPALIVE = 0x10  # macOS constant
        sock.setsockopt(socket.IPPROTO_TCP, TCP_KEEPALIVE, idle)
    elif sys.platform == "win32":
        # Windows: use SIO_KEEPALIVE_VALS ioctl
        try:
            import struct
            sock.ioctl(
                socket.SIO_KEEPALIVE_VALS,
                struct.pack("III", 1, idle * 1000, interval * 1000),
            )
        except (AttributeError, OSError):
            pass
    else:
        # Linux
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPIDLE, idle)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPINTVL, interval)
        sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_KEEPCNT, count)


# ------------------------------------------------------------------ #
# Connection pre-warming
# ------------------------------------------------------------------ #

def prewarm_connections(
    endpoints: list[tuple[str, int]] | None = None,
    timeout: float = 5.0,
) -> list[socket.socket]:
    """Establish and keep TCP connections to API endpoints.

    Returns list of connected sockets with keepalive enabled.
    These connections stay alive through CSCA's network blocking because:
    - Already-established TCP connections are not killed by firewall rule changes
    - Keepalive probes maintain the connection state
    """
    endpoints = endpoints or _WARMUP_ENDPOINTS
    sockets: list[socket.socket] = []

    for host, port in endpoints:
        try:
            sock = socket.create_connection((host, port), timeout=timeout)
            _set_tcp_keepalive(sock, idle=15, interval=5, count=3)
            sock.setblocking(False)
            sockets.append(sock)
            log.debug("Pre-warmed connection to %s:%d", host, port)
        except Exception as exc:
            log.debug("Failed to pre-warm %s:%d: %s", host, port, exc)

    if sockets:
        log.info("Pre-warmed %d/%d connections", len(sockets), len(endpoints))
    return sockets


# ------------------------------------------------------------------ #
# Proxy health watchdog
# ------------------------------------------------------------------ #

@dataclass
class NetworkProtector:
    """Keeps network connections alive through CSCA blocking.

    Features:
    - Pre-warms TCP connections to API endpoints before CSCA blocks
    - Monitors proxy availability with periodic health checks
    - Auto-reconnects if proxy connection is lost
    - Optionally adds WFP firewall rule to whitelist our process
    """

    proxy_url: str = ""
    check_interval: float = 30.0  # seconds between health checks
    enable_wfp_rule: bool = True  # try to add firewall rule (needs admin)
    enable_keepalive_patch: bool = True  # patch httpx to use keepalive

    _warmed_sockets: list[socket.socket] = field(default_factory=list, repr=False, init=False)
    _watchdog_thread: threading.Thread | None = field(default=None, repr=False, init=False)
    _stop_event: threading.Event = field(default_factory=threading.Event, repr=False, init=False)
    _ready_event: threading.Event = field(default_factory=threading.Event, repr=False, init=False)
    _wfp_rule_added: bool = field(default=False, repr=False, init=False)
    _proxy_alive: bool = field(default=False, repr=False, init=False)

    def start(self) -> None:
        """Start all network protections."""
        log.info("NetworkProtector starting...")

        # 1. DNS protection: warm cache + install override (before anything else)
        try:
            from .dns_protect import warm_dns_cache, install_dns_override

            warm_dns_cache()  # cache IPs before CSCA blocks DNS
            install_dns_override()  # monkey-patch socket.getaddrinfo
        except Exception as exc:
            log.debug("DNS protection init failed: %s", exc)

        # 2. Pre-warm connections to API endpoints
        self._warmed_sockets = prewarm_connections()

        # 3. Discover all available proxies (proxy pool)
        try:
            from .proxy_pool import ProxyPool

            self._proxy_pool = ProxyPool(check_interval=self.check_interval)
            count = self._proxy_pool.discover()
            if count > 0:
                best = self._proxy_pool.get_best()
                if best and not self.proxy_url:
                    self.proxy_url = best
                self._proxy_pool.start_monitor()
                log.info("Proxy pool: %d proxies discovered (best: %s)", count, best)
        except Exception as exc:
            self._proxy_pool = None
            log.debug("Proxy pool init failed: %s", exc)

        # 4. Start SOCKS5 proxy (fallback if no external proxy)
        # On Windows, try multiple bypass methods in order:
        #   a) SubprocessSocks5 — disguised python.exe (RuntimeBroker.exe etc.)
        #   b) PowerShellSocks5 — SOCKS5 in powershell.exe (CSCA trusts it)
        #   c) MicroSocks5 — in-process (works if CSCA doesn't block python.exe)
        if not self.proxy_url:
            self._socks5 = None
            self._socks5_sub = None
            self._socks5_ps = None

            if sys.platform == "win32":
                # Try a) Disguised subprocess
                try:
                    from .socks5_subprocess import SubprocessSocks5

                    self._socks5_sub = SubprocessSocks5(port=11081)
                    if self._socks5_sub.start():
                        self.proxy_url = self._socks5_sub.proxy_url
                        log.info("SOCKS5 subprocess proxy started: %s", self._socks5_sub.proxy_url)
                    else:
                        self._socks5_sub = None
                except Exception as exc:
                    self._socks5_sub = None
                    log.debug("SOCKS5 subprocess start failed: %s", exc)

                # Try b) PowerShell proxy (CSCA trusts powershell.exe)
                if not self.proxy_url:
                    try:
                        from .powershell_proxy import PowerShellSocks5

                        self._socks5_ps = PowerShellSocks5(port=11082)
                        if self._socks5_ps.start():
                            self.proxy_url = self._socks5_ps.proxy_url
                            log.info("PowerShell SOCKS5 proxy started: %s", self._socks5_ps.proxy_url)
                        else:
                            self._socks5_ps = None
                    except Exception as exc:
                        self._socks5_ps = None
                        log.debug("PowerShell SOCKS5 start failed: %s", exc)

            # Try c) In-process SOCKS5 (macOS/Linux, or if all above fail)
            if not self.proxy_url:
                try:
                    from .socks5_server import MicroSocks5

                    self._socks5 = MicroSocks5(listen_port=11081)
                    if self._socks5.start():
                        self.proxy_url = self._socks5.proxy_url
                        log.info("Built-in SOCKS5 proxy started: %s", self._socks5.proxy_url)
                    else:
                        self._socks5 = None
                except Exception as exc:
                    self._socks5 = None
                    log.debug("Built-in SOCKS5 start failed: %s", exc)
        else:
            self._socks5 = None
            self._socks5_sub = None
            self._socks5_ps = None

        # 5. Verify proxy is reachable
        if self.proxy_url:
            self._proxy_alive = self._check_proxy()
            if self._proxy_alive:
                log.info("Proxy verified: %s", self.proxy_url)
            else:
                log.warning("Proxy unreachable at start: %s", self.proxy_url)

        # 6. Add WFP firewall rule (Windows, requires admin)
        if self.enable_wfp_rule and sys.platform == "win32":
            self._wfp_rule_added = _add_wfp_allow_rule()
            # Also add rule for the disguised SOCKS5 subprocess
            if getattr(self, "_socks5_sub", None) and self._socks5_sub._disguised_exe:
                _add_wfp_allow_rule(exe_path=self._socks5_sub._disguised_exe)

        # 7. Patch httpx default transport to use keepalive
        if self.enable_keepalive_patch:
            _patch_httpx_keepalive()

        # 8. Start heartbeat connection to license server
        try:
            from .heartbeat import Heartbeat

            self._heartbeat = Heartbeat(host="license.muravey.app")
            self._heartbeat.start()
        except Exception as exc:
            self._heartbeat = None
            log.debug("Heartbeat init failed: %s", exc)

        # 9. Start watchdog thread
        self._stop_event.clear()
        self._watchdog_thread = threading.Thread(
            target=self._watchdog_loop,
            daemon=True,
            name="net-protect",
        )
        self._watchdog_thread.start()
        self._ready_event.set()
        log.info("NetworkProtector active (watchdog interval=%.0fs)", self.check_interval)

    def stop(self) -> None:
        """Stop all network protections and clean up."""
        self._stop_event.set()
        if self._watchdog_thread:
            self._watchdog_thread.join(timeout=5)
            self._watchdog_thread = None

        # Stop heartbeat
        if getattr(self, "_heartbeat", None):
            try:
                self._heartbeat.stop()
            except Exception:
                pass
            self._heartbeat = None

        # Stop proxy pool monitor
        if getattr(self, "_proxy_pool", None):
            try:
                self._proxy_pool.stop_monitor()
            except Exception:
                pass
            self._proxy_pool = None

        # Stop SOCKS5 subprocess
        if getattr(self, "_socks5_sub", None):
            try:
                self._socks5_sub.stop()
            except Exception:
                pass
            self._socks5_sub = None

        # Stop PowerShell SOCKS5
        if getattr(self, "_socks5_ps", None):
            try:
                self._socks5_ps.stop()
            except Exception:
                pass
            self._socks5_ps = None

        # Stop built-in SOCKS5
        if getattr(self, "_socks5", None):
            try:
                self._socks5.stop()
            except Exception:
                pass
            self._socks5 = None

        # Close pre-warmed sockets
        for sock in self._warmed_sockets:
            try:
                sock.close()
            except Exception:
                pass
        self._warmed_sockets.clear()

        # Remove WFP rule
        if self._wfp_rule_added:
            _remove_wfp_allow_rule()
            self._wfp_rule_added = False

        log.info("NetworkProtector stopped")

    @property
    def is_ready(self) -> bool:
        return self._ready_event.is_set()

    def wait_ready(self, timeout: float = 30.0) -> bool:
        """Block until NetworkProtector is fully started. Returns True if ready."""
        return self._ready_event.wait(timeout=timeout)

    @property
    def is_proxy_alive(self) -> bool:
        return self._proxy_alive

    def _check_proxy(self) -> bool:
        """Check if the SOCKS5 proxy can actually forward traffic."""
        if not self.proxy_url:
            return False

        from .proxy_tunnel import _parse_proxy_url
        parsed = _parse_proxy_url(self.proxy_url)
        if not parsed or not parsed.is_valid:
            return False

        # First check: port is open
        try:
            sock = socket.create_connection(
                (parsed.host, parsed.port), timeout=3.0,
            )
            sock.close()
        except (OSError, socket.timeout):
            return False

        # Second check: SOCKS5 handshake + connect through proxy
        # This verifies the proxy can actually forward, not just accept TCP
        try:
            import struct
            sock = socket.create_connection(
                (parsed.host, parsed.port), timeout=5.0,
            )
            # SOCKS5 greeting: version=5, 1 auth method, no-auth
            sock.sendall(b"\x05\x01\x00")
            resp = sock.recv(2)
            if len(resp) < 2 or resp[0] != 0x05 or resp[1] != 0x00:
                sock.close()
                log.warning("SOCKS5 handshake failed — proxy not usable")
                return False

            # SOCKS5 connect to dashscope.aliyuncs.com:443
            domain = b"dashscope.aliyuncs.com"
            req = struct.pack("!BBBB B", 0x05, 0x01, 0x00, 0x03, len(domain))
            req += domain + struct.pack("!H", 443)
            sock.sendall(req)
            resp = sock.recv(10)
            sock.close()

            if len(resp) >= 2 and resp[0] == 0x05 and resp[1] == 0x00:
                log.info("SOCKS5 proxy verified: can forward to internet")
                return True
            else:
                log.warning("SOCKS5 connect failed (rep=%s) — proxy not usable",
                            resp[1] if len(resp) >= 2 else "?")
                return False
        except Exception as exc:
            log.warning("SOCKS5 proxy forwarding check failed: %s", exc)
            return False

    def _watchdog_loop(self) -> None:
        """Periodic health check and auto-recovery."""
        while not self._stop_event.is_set():
            self._stop_event.wait(self.check_interval)
            if self._stop_event.is_set():
                break

            # Check proxy health
            if self.proxy_url:
                alive = self._check_proxy()
                if alive and not self._proxy_alive:
                    log.info("Proxy recovered: %s", self.proxy_url)
                elif not alive and self._proxy_alive:
                    log.warning("Proxy lost: %s — will retry", self.proxy_url)
                    # Try to re-detect proxy
                    self._try_rediscover_proxy()
                self._proxy_alive = alive

            # Re-warm any dropped connections
            self._refresh_warmed_sockets()

    def _try_rediscover_proxy(self) -> None:
        """Try to find a new proxy if the current one is dead."""
        # Try proxy pool first (instant failover)
        pool = getattr(self, "_proxy_pool", None)
        if pool:
            pool.mark_failed(self.proxy_url)
            best = pool.get_best()
            if best and best != self.proxy_url:
                log.info("Proxy pool failover: %s → %s", self.proxy_url, best)
                self.proxy_url = best
                self._proxy_alive = True
                return

        # Fallback: full re-scan
        try:
            from .proxy_tunnel import get_proxy_config, clear_proxy_cache

            clear_proxy_cache()
            config = get_proxy_config(force=True)
            if config and config.is_valid and config.url != self.proxy_url:
                log.info("Discovered new proxy: %s (was: %s)", config.url, self.proxy_url)
                self.proxy_url = config.url
                self._proxy_alive = True
        except Exception as exc:
            log.debug("Proxy re-discovery failed: %s", exc)

    def _refresh_warmed_sockets(self) -> None:
        """Check and re-establish pre-warmed connections."""
        alive = []
        for sock in self._warmed_sockets:
            try:
                # Non-blocking recv to check if socket is still connected
                sock.setblocking(False)
                data = sock.recv(1, socket.MSG_PEEK)
                if data == b"":
                    sock.close()
                    continue
                alive.append(sock)
            except BlockingIOError:
                # No data available — socket is still alive
                alive.append(sock)
            except (OSError, socket.timeout):
                try:
                    sock.close()
                except Exception:
                    pass

        dropped = len(self._warmed_sockets) - len(alive)
        if dropped > 0:
            log.debug("Re-warming %d dropped connections", dropped)
            new_socks = prewarm_connections()
            alive.extend(new_socks)

        self._warmed_sockets = alive


# ------------------------------------------------------------------ #
# WFP firewall rule (Windows only, requires admin)
# ------------------------------------------------------------------ #

_WFP_RULE_NAME = "MuraveyChina Network Allow"


def _add_wfp_allow_rule(exe_path: str | None = None) -> bool:
    """Add Windows Firewall rule to allow a process outbound traffic.

    This ensures CSCA's firewall rules don't block us.
    Requires admin privileges — fails silently if not admin.
    """
    if sys.platform != "win32":
        return False

    try:
        # Check if we have admin rights
        is_admin = ctypes.windll.shell32.IsUserAnAdmin() != 0
        if not is_admin:
            log.debug("WFP rule skipped: not running as admin")
            return False
    except Exception:
        return False

    exe_path = exe_path or sys.executable
    try:
        # Add outbound allow rule for our process
        cmd = [
            "netsh", "advfirewall", "firewall", "add", "rule",
            f"name={_WFP_RULE_NAME}",
            "dir=out",
            "action=allow",
            f"program={exe_path}",
            "enable=yes",
            "profile=any",
        ]
        subprocess.run(
            cmd, capture_output=True, timeout=10,
            creationflags=subprocess.CREATE_NO_WINDOW if hasattr(subprocess, "CREATE_NO_WINDOW") else 0,
        )

        # Also add inbound rule for SSH tunnel / proxy
        cmd_in = [
            "netsh", "advfirewall", "firewall", "add", "rule",
            f"name={_WFP_RULE_NAME} (in)",
            "dir=in",
            "action=allow",
            f"program={exe_path}",
            "enable=yes",
            "profile=any",
            "localport=11080,7890,1080",
            "protocol=tcp",
        ]
        subprocess.run(
            cmd_in, capture_output=True, timeout=10,
            creationflags=subprocess.CREATE_NO_WINDOW if hasattr(subprocess, "CREATE_NO_WINDOW") else 0,
        )

        log.info("WFP firewall rule added: allow %s", os.path.basename(exe_path))
        return True
    except Exception as exc:
        log.debug("WFP rule creation failed: %s", exc)
        return False


def _remove_wfp_allow_rule() -> bool:
    """Remove our firewall rules on cleanup."""
    if sys.platform != "win32":
        return False

    try:
        for name in [_WFP_RULE_NAME, f"{_WFP_RULE_NAME} (in)"]:
            subprocess.run(
                ["netsh", "advfirewall", "firewall", "delete", "rule", f"name={name}"],
                capture_output=True, timeout=10,
                creationflags=subprocess.CREATE_NO_WINDOW if hasattr(subprocess, "CREATE_NO_WINDOW") else 0,
            )
        return True
    except Exception:
        return False


# ------------------------------------------------------------------ #
# httpx keepalive patch
# ------------------------------------------------------------------ #

_httpx_patched = False


def _patch_httpx_keepalive() -> None:
    """Patch httpx default transport to use aggressive keepalive.

    This ensures all HTTP connections made through httpx (including OpenAI SDK)
    use TCP keepalive to survive CSCA's idle connection killing.
    """
    global _httpx_patched
    if _httpx_patched:
        return

    try:
        import httpx
        import httpcore

        # Patch default pool keepalive settings
        _orig_init = httpcore.ConnectionPool.__init__

        def _patched_init(self, *args, **kwargs):
            # More aggressive keepalive: 60s expiry instead of 5s
            kwargs.setdefault("keepalive_expiry", 60.0)
            _orig_init(self, *args, **kwargs)

        httpcore.ConnectionPool.__init__ = _patched_init
        _httpx_patched = True
        log.debug("httpx keepalive patched (60s expiry)")
    except Exception as exc:
        log.debug("httpx keepalive patch skipped: %s", exc)


# ------------------------------------------------------------------ #
# High-level: init network protection for stealth mode
# ------------------------------------------------------------------ #

_active_protector: NetworkProtector | None = None


def init_network_protection(proxy_url: str = "", check_interval: float = 30.0) -> NetworkProtector:
    """Create and start a NetworkProtector instance.

    Call this BEFORE CSCA starts to pre-warm connections.
    """
    global _active_protector
    protector = NetworkProtector(
        proxy_url=proxy_url,
        check_interval=check_interval,
    )
    protector.start()
    _active_protector = protector
    return protector


def wait_for_network(timeout: float = 30.0) -> bool:
    """Wait until NetworkProtector is ready. Returns True if ready, False on timeout.

    If no NetworkProtector was started, returns True immediately (no protection needed).
    """
    if _active_protector is None:
        return True
    return _active_protector.wait_ready(timeout=timeout)
