"""DNS protection — DoH resolution and IP pinning for CSCA bypass.

When CSCA blocks DNS queries, this module provides:
1. DNS over HTTPS (DoH) via Cloudflare/Google — bypasses DNS blocking
2. DNS pinning — caches IPs at startup, uses cached IPs when DNS is blocked

Usage:
    from src.utils.dns_protect import resolve_host, warm_dns_cache

    # Pre-cache IPs at startup (before CSCA blocks DNS)
    warm_dns_cache()

    # Resolve with fallback: regular DNS → DoH → cached IP
    ip = resolve_host("api.deepseek.com")
"""

from __future__ import annotations

import json
import logging
import socket
import ssl
import threading
import time
import urllib.request
from dataclasses import dataclass, field

log = logging.getLogger(__name__)

# Domains critical for MuraveyChina operation
CRITICAL_DOMAINS: list[str] = [
    "api.deepseek.com",
    "dashscope.aliyuncs.com",
    "license.muravey.app",
    "api.88api.dev",
]

# DoH endpoints (Cloudflare primary, Google fallback)
_DOH_ENDPOINTS: list[tuple[str, str]] = [
    ("1.1.1.1", "https://1.1.1.1/dns-query"),
    ("8.8.8.8", "https://dns.google/resolve"),
    ("223.5.5.5", "https://dns.alidns.com/resolve"),  # Alibaba DNS (works in China)
]


@dataclass
class DNSCache:
    """Thread-safe DNS cache with TTL."""

    _cache: dict[str, tuple[str, float]] = field(default_factory=dict)  # domain -> (ip, expiry_ts)
    _lock: threading.Lock = field(default_factory=threading.Lock)
    default_ttl: float = 3600.0  # 1 hour

    def get(self, domain: str) -> str | None:
        with self._lock:
            entry = self._cache.get(domain)
            if entry is None:
                return None
            ip, expiry = entry
            # Return even if expired (stale cache is better than no DNS)
            return ip

    def put(self, domain: str, ip: str, ttl: float | None = None) -> None:
        with self._lock:
            self._cache[domain] = (ip, time.monotonic() + (ttl or self.default_ttl))

    def get_all(self) -> dict[str, str]:
        with self._lock:
            return {domain: ip for domain, (ip, _) in self._cache.items()}

    @property
    def size(self) -> int:
        with self._lock:
            return len(self._cache)


_dns_cache = DNSCache()


# Regular DNS resolution
def _resolve_regular(domain: str, timeout: float = 3.0) -> str | None:
    try:
        old_timeout = socket.getdefaulttimeout()
        socket.setdefaulttimeout(timeout)
        try:
            result = socket.getaddrinfo(domain, 443, socket.AF_INET)
            if result:
                return result[0][4][0]
        finally:
            socket.setdefaulttimeout(old_timeout)
    except (socket.gaierror, socket.timeout, OSError):
        pass
    return None


# DoH resolution via Cloudflare JSON API
def _resolve_doh_cloudflare(domain: str, doh_ip: str, doh_url: str, timeout: float = 5.0) -> str | None:
    """Resolve domain using DNS over HTTPS (JSON API)."""
    try:
        from .proxied_urllib import proxied_urlopen

        url = f"{doh_url}?name={domain}&type=A"
        req = urllib.request.Request(url, headers={
            "Accept": "application/dns-json",
        })

        # Create SSL context that doesn't verify (since we connect by IP)
        ctx = ssl.create_default_context()

        with proxied_urlopen(req, timeout=timeout) as resp:
            data = json.loads(resp.read())

        # Parse response (Cloudflare and Google use similar JSON format)
        answers = data.get("Answer", [])
        for answer in answers:
            if answer.get("type") == 1:  # A record
                ip = answer.get("data", "")
                if ip:
                    return ip
    except Exception as exc:
        log.debug("DoH resolution via %s failed for %s: %s", doh_ip, domain, exc)
    return None


def resolve_host(domain: str, timeout: float = 3.0) -> str | None:
    """Resolve domain with fallback chain: regular DNS → DoH → cached IP.

    Returns IP address string or None if all methods fail.
    """
    # 1. Try regular DNS
    ip = _resolve_regular(domain, timeout=timeout)
    if ip:
        _dns_cache.put(domain, ip)
        return ip

    log.debug("Regular DNS failed for %s, trying DoH...", domain)

    # 2. Try DoH endpoints
    for doh_ip, doh_url in _DOH_ENDPOINTS:
        ip = _resolve_doh_cloudflare(domain, doh_ip, doh_url, timeout=timeout + 2)
        if ip:
            _dns_cache.put(domain, ip)
            log.info("DoH resolved %s → %s (via %s)", domain, ip, doh_ip)
            return ip

    log.debug("DoH failed for %s, checking cache...", domain)

    # 3. Fall back to cached IP (stale is better than nothing)
    cached = _dns_cache.get(domain)
    if cached:
        log.info("Using cached IP for %s → %s (DNS blocked)", domain, cached)
        return cached

    log.warning("All DNS resolution failed for %s", domain)
    return None


def warm_dns_cache(domains: list[str] | None = None) -> dict[str, str]:
    """Pre-resolve and cache IPs for critical domains.

    Call this at startup BEFORE CSCA blocks DNS.
    Returns dict of domain → IP for successfully resolved domains.
    """
    domains = domains or CRITICAL_DOMAINS
    resolved: dict[str, str] = {}

    for domain in domains:
        ip = _resolve_regular(domain, timeout=3.0)
        if ip:
            _dns_cache.put(domain, ip, ttl=86400.0)  # 24h TTL for warm cache
            resolved[domain] = ip
        else:
            # Try DoH immediately for domains that fail regular DNS
            for doh_ip, doh_url in _DOH_ENDPOINTS:
                ip = _resolve_doh_cloudflare(domain, doh_ip, doh_url)
                if ip:
                    _dns_cache.put(domain, ip, ttl=86400.0)
                    resolved[domain] = ip
                    break

    log.info("DNS cache warmed: %d/%d domains resolved", len(resolved), len(domains))
    return resolved


def warm_dns_cache_async(domains: list[str] | None = None) -> threading.Thread:
    """Non-blocking version of warm_dns_cache."""
    t = threading.Thread(target=warm_dns_cache, args=(domains,), daemon=True, name="dns-warm")
    t.start()
    return t


def get_cached_ip(domain: str) -> str | None:
    """Get cached IP without triggering resolution."""
    return _dns_cache.get(domain)


def get_all_cached() -> dict[str, str]:
    """Get all cached domain → IP mappings."""
    return _dns_cache.get_all()


def install_dns_override() -> None:
    """Monkey-patch socket.getaddrinfo to use our DNS cache as fallback.

    This ensures ALL socket connections (including OpenAI SDK, httpx, etc.)
    benefit from our DoH + cache fallback without code changes.
    """
    _original_getaddrinfo = socket.getaddrinfo

    def _patched_getaddrinfo(host, port, family=0, type=0, proto=0, flags=0):
        try:
            return _original_getaddrinfo(host, port, family, type, proto, flags)
        except socket.gaierror:
            # Regular DNS failed — try our resolution
            ip = resolve_host(host)
            if ip:
                log.debug("DNS override: %s → %s", host, ip)
                # Return a valid addrinfo tuple
                return [(socket.AF_INET, socket.SOCK_STREAM, 6, '', (ip, port or 443))]
            raise  # Re-raise if we also can't resolve

    socket.getaddrinfo = _patched_getaddrinfo
    log.info("DNS override installed (DoH + cache fallback)")
