from __future__ import annotations

import asyncio
import base64
import hashlib
import json
import os
import random
import re
import socket
import time
from dataclasses import dataclass
from typing import Any, Dict, List, Mapping, Optional

import httpx
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes

from .utils import normalize_ip


INTEL_KEYS = [
    "banned_isps",
    "banned_asns",
    "scanner_uas",
    "proxy_domains",
    "cloudflare_cidrs",
    "bot_ua_regex",
    "banned_referrers",
]


class UltimateProtectorClient:
    agent_kind = "python"
    agent_version = "0.1.0"
    protocol_version = 1
    capabilities = ["enforce:block", "enforce:challenge", "telemetry:v1", "obsidian"]

    max_ua_length = 1024

    def __init__(
        self,
        *,
        license_key: str,
        api_url: str,
        sync_interval_seconds: int,
        allow_sample_rate: float,
    ) -> None:
        self.license_key = str(license_key)
        self.api_url = str(api_url).rstrip("/")
        self.sync_interval_seconds = max(10, int(sync_interval_seconds))
        self.allow_sample_rate = max(0.0, min(1.0, float(allow_sample_rate)))

        self.status = "empty"  # empty|success|expired
        self.synced_at = 0
        self.rules: Optional[Dict[str, Any]] = None

        safe = hashlib.sha256(self.license_key.encode("utf-8")).hexdigest()[:16]
        self.cache_path = os.path.join(temp_dir(), f"up_rules_{safe}.json")

        self._refresh_task: Optional[asyncio.Task] = None
        self._refresh_lock = asyncio.Lock()

        self._blocked_ip_set = None
        self._blocked_ip_set_version = None

        # RDNS cache
        self.rdns_ttl_seconds = 86400
        self.rdns_max_entries = 5000
        self._rdns_cache: Dict[str, Dict[str, Any]] = {}

    def expired_html(self) -> str:
        return (
            "<!DOCTYPE html><html lang=\"en\"><head><meta charset=\"UTF-8\" />"
            "<meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\" />"
            "<title>Service Notice</title>"
            "<style>body{background:#050507;color:#fff;font-family:sans-serif;display:flex;align-items:center;justify-content:center;height:100vh;margin:0}.card{background:#0E0E10;padding:40px;border-radius:20px;border:1px solid rgba(255,255,255,.1);text-align:center}h1{color:#ef4444}</style>"
            "</head><body><div class=\"card\"><h1>Subscription Expired</h1><p>Security license inactive. Protection disabled.</p>"
            "<script>setTimeout(()=>location.reload(),6000)</script></div></body></html>"
        )

    def verify_up_token(self, token: str) -> bool:
        try:
            decoded = base64.b64decode(token.encode("utf-8"), validate=False).decode("utf-8")
            if "::" not in decoded:
                return False
            ts_raw, sig = decoded.split("::", 1)
            ts = int(ts_raw)
            if time.time() - ts > 60:
                return False
            raw = f"{ts}|{self.license_key}".encode("utf-8")
            calc = hmac_sha256(self.license_key.encode("utf-8"), raw)
            return safe_equals(calc, sig.encode("utf-8"))
        except Exception:
            return False

    async def get_rules(self, *, domain: str) -> Optional[Dict[str, Any]]:
        if self.status == "empty":
            self._load_from_disk()

        now = int(time.time())
        stale = (now - int(self.synced_at or 0)) > self.sync_interval_seconds

        if self.rules is None and self.status != "expired":
            await self.refresh_rules(domain=domain)
            return self.rules

        if stale and self.status != "expired":
            self._schedule_refresh(domain=domain)

        return self.rules

    def _schedule_refresh(self, *, domain: str) -> None:
        if self._refresh_task and not self._refresh_task.done():
            return
        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            return
        self._refresh_task = loop.create_task(self.refresh_rules(domain=domain))

    async def refresh_rules(self, *, domain: str) -> None:
        async with self._refresh_lock:
            blocked_ver = int(self.rules.get("blocked_ips_version") or 0) if isinstance(self.rules, dict) else 0
            intel_ver = int(self.rules.get("intel_version") or 0) if isinstance(self.rules, dict) else 0

            payload = {
                "license_key": self.license_key,
                "domain": domain or "unknown",
                "blocked_ips_version": blocked_ver,
                "intel_version": intel_ver,
                "protocol_version": self.protocol_version,
                "agent_kind": self.agent_kind,
                "agent_version": self.agent_version,
                "capabilities": self.capabilities,
            }

            try:
                async with httpx.AsyncClient(timeout=5.0) as client:
                    r = await client.post(self.api_url + "/verify", json=payload)
                    data = r.json()
            except Exception:
                return  # fail-open

            now = int(time.time())

            if isinstance(data, dict) and data.get("status") == "expired":
                self.status = "expired"
                self.rules = None
                self.synced_at = now
                self._save_to_disk()
                return

            if isinstance(data, dict) and data.get("status") == "success":
                if data.get("encrypted") is True and data.get("payload"):
                    decoded = self._decrypt_rules_payload(str(data["payload"]))
                    if isinstance(decoded, dict):
                        merged = dict(decoded)
                        if isinstance(self.rules, dict):
                            for k in INTEL_KEYS:
                                if k in merged and merged.get(k) is None:
                                    merged[k] = self.rules.get(k)
                            if "global_blocked_ips" in merged and merged.get("global_blocked_ips") is None:
                                merged["global_blocked_ips"] = self.rules.get("global_blocked_ips")

                        self.status = "success"
                        self.rules = merged
                        self.synced_at = now
                        self._save_to_disk()
                        return

                # success but empty payload (domain blocked/revoked)
                if data.get("encrypted") is False and (data.get("payload") is None):
                    self.status = "success"
                    self.rules = None
                    self.synced_at = now
                    self._save_to_disk()
                    return

    def _decrypt_rules_payload(self, b64: str) -> Optional[Dict[str, Any]]:
        try:
            raw = base64.b64decode(b64)
            if len(raw) < 17:
                return None
            iv = raw[:16]
            ciphertext = raw[16:]
            key = hashlib.sha256(self.license_key.encode("utf-8")).digest()

            cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
            decryptor = cipher.decryptor()
            padded = decryptor.update(ciphertext) + decryptor.finalize()
            plain = pkcs7_unpad(padded)
            data = json.loads(plain.decode("utf-8"))
            return data if isinstance(data, dict) else None
        except Exception:
            return None

    def _load_from_disk(self) -> None:
        try:
            with open(self.cache_path, "r", encoding="utf-8") as f:
                data = json.load(f)
            if not isinstance(data, dict):
                return
            if data.get("status") == "expired":
                self.status = "expired"
                self.synced_at = int(data.get("synced_at") or 0)
                self.rules = None
                return
            if isinstance(data.get("rules"), dict):
                self.status = "success"
                self.synced_at = int(data.get("synced_at") or 0)
                self.rules = data.get("rules")
        except Exception:
            return

    def _save_to_disk(self) -> None:
        try:
            payload = {"synced_at": int(self.synced_at), "status": self.status}
            if self.status == "success":
                payload["rules"] = self.rules
            tmp = self.cache_path + ".tmp"
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(payload, f)
            os.replace(tmp, self.cache_path)
        except Exception:
            return

    def is_globally_blocked_ip(self, rules: Mapping[str, Any], ip: str) -> bool:
        ver = rules.get("blocked_ips_version")
        if ver is None:
            return False
        if self._blocked_ip_set is None or self._blocked_ip_set_version != ver:
            ips = rules.get("global_blocked_ips")
            self._blocked_ip_set = set(ips) if isinstance(ips, list) else set()
            self._blocked_ip_set_version = ver
        return ip in self._blocked_ip_set

    async def log_allow_verified_human(self, *, ip: str, ua: str, method: str, host: str, path: str, country: str) -> None:
        await self.log_allow("Verified Human", None, {
            "ip": ip,
            "ua": ua,
            "method": method,
            "host": host,
            "path": path,
            "country": country,
        })

    async def log_allow(self, reason: str, reason_code: Optional[str], ctx: Mapping[str, Any]) -> None:
        await self._log("ALLOW", reason, reason_code, ctx)

    async def log_block(self, reason: str, reason_code: Optional[str], ctx: Mapping[str, Any]) -> None:
        await self._log("BLOCK", reason, reason_code, ctx, force=True)

    async def log_challenge(self, reason: str, reason_code: Optional[str], ctx: Mapping[str, Any]) -> None:
        await self._log("CHALLENGE", reason, reason_code, ctx)

    async def _log(self, action: str, reason: str, reason_code: Optional[str], ctx: Mapping[str, Any], *, force: bool = False) -> None:
        ip = normalize_ip(ctx.get("ip"))
        if not ip:
            return

        if action == "ALLOW" and (not force):
            if self.allow_sample_rate <= 0.0:
                return
            if random.random() > self.allow_sample_rate:
                return

        body: Dict[str, Any] = {
            "license_key": self.license_key,
            "ip": ip,
            "action": action,
            "reason": reason,
            "reason_code": reason_code,
            "user_agent": ctx.get("ua") or "",
            "ua_hash": hashlib.sha256((ctx.get("ua") or "").encode("utf-8")).hexdigest() if ctx.get("ua") else None,
            "country": ctx.get("country") or "XX",
            "method": ctx.get("method") or "",
            "host": ctx.get("host") or "",
            "path": ctx.get("path") or "/",
            "protocol_version": self.protocol_version,
            "agent_kind": self.agent_kind,
            "agent_version": self.agent_version,
            "capabilities": self.capabilities,
        }

        # Fire-and-forget
        try:
            asyncio.create_task(self._post_log(body))
        except RuntimeError:
            # no running loop
            return

    async def _post_log(self, body: Dict[str, Any]) -> None:
        try:
            async with httpx.AsyncClient(timeout=2.0) as client:
                await client.post(self.api_url + "/log", json=body)
        except Exception:
            return

    async def rdns_lookup(self, ip: str) -> Optional[str]:
        ip = normalize_ip(ip)
        if not ip:
            return None

        now = int(time.time())
        hit = self._rdns_cache.get(ip)
        if isinstance(hit, dict) and (now - int(hit.get("t") or 0)) < self.rdns_ttl_seconds:
            return hit.get("h")

        hostname = await asyncio.to_thread(self._dns_reverse, ip)
        self._rdns_cache[ip] = {**(hit or {}), "t": now, "h": hostname}

        if len(self._rdns_cache) > self.rdns_max_entries:
            items = sorted(self._rdns_cache.items(), key=lambda kv: int((kv[1] or {}).get("t") or 0))
            drop = max(1, int(self.rdns_max_entries * 0.1))
            for k, _ in items[:drop]:
                self._rdns_cache.pop(k, None)

        return hostname

    async def is_safe_seo_crawler(self, *, ua: str, ip: str) -> bool:
        ua_l = (ua or "").lower()
        if any(t in ua_l for t in [
            "googlebot",
            "adsbot-google",
            "mediapartners-google",
            "google-inspectiontool",
            "googleother",
            "google-extended",
            "apis-google",
        ]):
            return await self._is_verified_google_ip(ip)

        if not any(t in ua_l for t in [
            "bingbot", "msnbot",
            "yandex",
            "baiduspider",
            "duckduckbot",
            "applebot",
            "sogou",
            "exabot",
            "seznambot",
            "naverbot",
            "petalbot",
            "bytespider",
            "slurp",
        ]):
            return False

        return await self._is_verified_seo_ip(ip)

    async def _is_verified_google_ip(self, ip: str) -> bool:
        ip = normalize_ip(ip)
        if not ip:
            return False

        now = int(time.time())
        hit = self._rdns_cache.get(ip)
        if isinstance(hit, dict) and "gv" in hit and (now - int(hit.get("t") or 0)) < self.rdns_ttl_seconds:
            return bool(hit.get("gv"))

        host = await self.rdns_lookup(ip)
        h = (host or "").lower().rstrip(".")
        verified = False
        if h and (h.endswith(".googlebot.com") or h.endswith(".google.com")):
            ips = await asyncio.to_thread(self._dns_forward_ips, h)
            verified = ip in ips

        self._rdns_cache[ip] = {**(hit or {}), "t": now, "h": host, "gv": verified}
        return verified

    async def _is_verified_seo_ip(self, ip: str) -> bool:
        ip = normalize_ip(ip)
        if not ip:
            return False

        suffixes = [
            ".search.msn.com",
            ".yandex.ru", ".yandex.net", ".yandex.com",
            ".baidu.com",
            ".duckduckgo.com",
            ".applebot.apple.com",
            ".sogou.com",
            ".exabot.com",
            ".seznam.cz",
            ".naver.com",
            ".petalbot.com",
            ".bytespider.com",
            ".crawl.yahoo.net", ".yahoo.com",
        ]

        now = int(time.time())
        hit = self._rdns_cache.get(ip)
        if isinstance(hit, dict) and "sv" in hit and (now - int(hit.get("t") or 0)) < self.rdns_ttl_seconds:
            return bool(hit.get("sv"))

        host = await self.rdns_lookup(ip)
        h = (host or "").lower().rstrip(".")
        verified = False
        if h:
            for s in suffixes:
                if h.endswith(s):
                    ips = await asyncio.to_thread(self._dns_forward_ips, h)
                    verified = ip in ips
                    break

        self._rdns_cache[ip] = {**(hit or {}), "t": now, "h": host, "sv": verified}
        return verified

    @staticmethod
    def _dns_reverse(ip: str) -> Optional[str]:
        try:
            host, _, _ = socket.gethostbyaddr(ip)
            return host if host and host != ip else None
        except Exception:
            return None

    @staticmethod
    def _dns_forward_ips(hostname: str) -> List[str]:
        host = (hostname or "").strip().lower().rstrip(".")
        if not host:
            return []
        ips: List[str] = []
        try:
            infos = socket.getaddrinfo(host, None)
            for fam, _, _, _, sockaddr in infos:
                if fam in (socket.AF_INET, socket.AF_INET6):
                    ip = normalize_ip(sockaddr[0])
                    if ip:
                        ips.append(ip)
        except Exception:
            return []
        # unique
        return sorted(set(ips))


def temp_dir() -> str:
    return os.getenv("TMPDIR") or os.getenv("TEMP") or os.getenv("TMP") or "/tmp"


def pkcs7_unpad(data: bytes) -> bytes:
    if not data:
        raise ValueError("empty")
    pad = data[-1]
    if pad < 1 or pad > 16:
        raise ValueError("bad padding")
    if data[-pad:] != bytes([pad]) * pad:
        raise ValueError("bad padding")
    return data[:-pad]


def hmac_sha256(key: bytes, msg: bytes) -> bytes:
    import hmac

    return hmac.new(key, msg, digestmod=hashlib.sha256).hexdigest().encode("utf-8")


def safe_equals(a: bytes, b: bytes) -> bool:
    try:
        import hmac

        return hmac.compare_digest(a, b)
    except Exception:
        return False
