from __future__ import annotations

import asyncio
import hashlib
import os
import time
from typing import Any, Awaitable, Callable, Dict, Iterable, List, Mapping, Optional, Tuple

from .client import UltimateProtectorClient
from .obsidian import ObsidianInjector
from .utils import (
    cidr_match,
    compile_php_like_regex,
    detect_country,
    is_sensitive_request,
    is_static_request,
    normalize_ip,
    parse_cookie_header,
    should_protect_request,
)

ASGIApp = Callable[[Dict[str, Any], Callable[[], Awaitable[Dict[str, Any]]], Callable[[Dict[str, Any]], Awaitable[None]]], Awaitable[None]]


class UltimateProtectorMiddleware:
    def __init__(
        self,
        app: ASGIApp,
        *,
        license_key: str,
        api_url: str,
        sync_interval_seconds: int = 60,
        allow_sample_rate: float = 0.01,
        only_paths: Optional[List[str]] = None,
        except_paths: Optional[List[str]] = None,
        only_regex: Optional[str] = None,
    ) -> None:
        self.app = app
        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.only_paths = [str(x) for x in only_paths] if only_paths else None
        self.except_paths = [str(x) for x in except_paths] if except_paths else None
        self.only_regex = str(only_regex) if only_regex else None

        self.client = UltimateProtectorClient(
            license_key=self.license_key,
            api_url=self.api_url,
            sync_interval_seconds=self.sync_interval_seconds,
            allow_sample_rate=self.allow_sample_rate,
        )

        self.obsidian = ObsidianInjector()

        self.passport_cookie_name = "aura_passport"
        self.passport_cookie_value = hashlib.sha256(("verified" + self.license_key).encode("utf-8")).hexdigest()

    async def __call__(self, scope: Dict[str, Any], receive: Callable[[], Awaitable[Dict[str, Any]]], send: Callable[[Dict[str, Any]], Awaitable[None]]):
        if scope.get("type") != "http":
            await self.app(scope, receive, send)
            return

        path = str(scope.get("path") or "/")
        if not should_protect_request(
            path=path,
            only_paths=self.only_paths,
            except_paths=self.except_paths,
            only_regex=self.only_regex,
        ):
            await self.app(scope, receive, send)
            return

        headers = _headers_to_dict(scope.get("headers") or [])
        method = str(scope.get("method") or "").upper()

        qs = scope.get("query_string") or b""
        query = _parse_query_string(qs)

        # up_token exchange (challenge success)
        up_token = query.get("up_token")
        if up_token:
            if not self.client.verify_up_token(up_token):
                await _send_text(send, 403, "Invalid Token")
                return

            secure = _is_secure(scope, headers)
            set_cookie = _build_cookie(
                name=self.passport_cookie_name,
                value=self.passport_cookie_value,
                max_age=86400,
                secure=secure,
            )

            await _send_redirect(send, 302, _absolute_url(scope, headers, path), extra_headers=[(b"set-cookie", set_cookie.encode("utf-8"))])
            return

        cookies = parse_cookie_header(headers.get("cookie", ""))
        if cookies.get(self.passport_cookie_name) == self.passport_cookie_value:
            ip = normalize_ip(_client_ip_from_scope(scope))
            ua = headers.get("user-agent", "")
            await self.client.log_allow_verified_human(
                ip=ip,
                ua=ua,
                method=method,
                host=headers.get("host", ""),
                path=path,
                country=detect_country(headers),
            )
            await self.app(scope, receive, send)
            return

        host = headers.get("host", "")
        domain = host.split(":")[0] if host else "unknown"

        try:
            rules = await self.client.get_rules(domain=domain)
        except Exception:
            # fail-open
            await self.app(scope, receive, send)
            return

        if self.client.status == "expired":
            await _send_html(send, 503, self.client.expired_html())
            return

        if not isinstance(rules, dict):
            await self.app(scope, receive, send)
            return

        # Obsidian hooks (response transform)
        send_wrapped = send
        if rules.get("obsidian_active"):
            send_wrapped = self.obsidian.wrap_send(send, rules)

        # Context
        ua = (headers.get("user-agent") or "")[: self.client.max_ua_length]
        country = detect_country(headers)
        origin = headers.get("origin", "")
        ref = headers.get("referer", "")

        remote_ip = normalize_ip(_client_ip_from_scope(scope))
        ip = _detect_client_ip(headers, rules, remote_ip)

        ctx = {
            "ip": ip,
            "ua": ua,
            "method": method,
            "host": host,
            "path": path,
            "country": country,
        }

        # L-1 whitelist
        wl = rules.get("whitelist_ips")
        if isinstance(wl, list) and ip in wl:
            await self.client.log_allow("Whitelisted IP", None, ctx)
            await self.app(scope, receive, send_wrapped)
            return

        # L-0.5 SEO safety
        if rules.get("seo_safety_enabled"):
            ok = await self.client.is_safe_seo_crawler(ua=ua, ip=ip)
            if ok:
                await self.client.log_allow("SEO Safety", None, ctx)
                await self.app(scope, receive, send_wrapped)
                return

        # L1 global blocklist
        if self.client.is_globally_blocked_ip(rules, ip):
            await self.client.log_block("Global Blocklist", "global_blocklist", ctx)
            await _respond_block(send_wrapped, rules, ip, "Global Blocklist")
            return

        # L2 geo firewall
        block_geo = rules.get("block_geo")
        if isinstance(block_geo, list) and block_geo:
            is_whitelist = (rules.get("geo_mode") or "blacklist") == "whitelist"
            in_list = country in block_geo
            if (is_whitelist and not in_list) or ((not is_whitelist) and in_list):
                await self.client.log_block("Geo Firewall", "geo", ctx)
                await _respond_block(send_wrapped, rules, ip, "Geo Firewall")
                return

        # L3 scanner UAs
        scanner_uas = rules.get("scanner_uas")
        if isinstance(scanner_uas, list):
            ua_lower = ua.lower()
            for bad in scanner_uas:
                if isinstance(bad, str) and bad and bad.lower() in ua_lower:
                    await self.client.log_block("Bot Signature", "bot_signature", ctx)
                    await _respond_block(send_wrapped, rules, ip, "Bot Signature")
                    return

        # L3 bot regex
        bot_ua_regex = rules.get("bot_ua_regex")
        if isinstance(bot_ua_regex, str) and bot_ua_regex:
            re_obj = compile_php_like_regex(bot_ua_regex)
            if re_obj and re_obj.search(ua):
                await self.client.log_block("Bot Regex", "bot_regex", ctx)
                await _respond_block(send_wrapped, rules, ip, "Bot Regex")
                return

        # L4 VPN shield (heuristics)
        if rules.get("block_vpn"):
            uri = _absolute_url(scope, headers, path)
            do_rdns = (not is_static_request(uri)) and method != "HEAD" and is_sensitive_request(method, uri)
            hostname = None
            if do_rdns:
                hostname = await self.client.rdns_lookup(ip)

            if hostname and hostname != ip:
                proxy_domains = rules.get("proxy_domains")
                if isinstance(proxy_domains, list):
                    hn = hostname.lower()
                    for d in proxy_domains:
                        if isinstance(d, str) and d and d.lower() in hn:
                            await self.client.log_challenge("Challenge", "challenge", ctx)
                            await _respond_challenge(send_wrapped, scope, headers, self.license_key, self.api_url)
                            return

                banned_isps = rules.get("banned_isps")
                if isinstance(banned_isps, list):
                    hn = hostname.lower()
                    for s in banned_isps:
                        if isinstance(s, str) and s and s.lower() in hn:
                            await self.client.log_block("Banned ISP", "banned_isp", ctx)
                            await _respond_block(send_wrapped, rules, ip, "Banned ISP")
                            return

                banned_asns = rules.get("banned_asns")
                if isinstance(banned_asns, list):
                    hn = hostname.lower()
                    for s in banned_asns:
                        if isinstance(s, str) and s and s.lower() in hn:
                            await self.client.log_block("Banned ASN", "banned_asn", ctx)
                            await _respond_block(send_wrapped, rules, ip, "Banned ASN")
                            return

        # L5 referrer security
        banned_ref = rules.get("banned_referrers")
        if isinstance(banned_ref, list) and banned_ref:
            if method == "POST" and not ref and not origin:
                await self.client.log_challenge("Challenge", "challenge", ctx)
                await _respond_challenge(send_wrapped, scope, headers, self.license_key, self.api_url)
                return

            if ref:
                # protocol mismatch detection (https site with http referrer)
                try:
                    from urllib.parse import urlparse

                    ru = urlparse(ref)
                    ref_host = ru.netloc
                    ref_scheme = ru.scheme
                    is_https = _is_secure(scope, headers)
                    if ref_host == host and is_https and ref_scheme == "http":
                        await self.client.log_block("Header Spoofing", None, ctx)
                        await _respond_block(send_wrapped, rules, ip, "Header Spoofing")
                        return
                except Exception:
                    pass

                ref_lower = ref.lower()
                for bad in banned_ref:
                    if isinstance(bad, str) and bad and bad.lower() in ref_lower:
                        await self.client.log_block("Bad Referrer", "bad_referrer", ctx)
                        await _respond_block(send_wrapped, rules, ip, "Bad Referrer")
                        return

        await self.client.log_allow("Clean Traffic", None, ctx)
        await self.app(scope, receive, send_wrapped)


def _headers_to_dict(headers: Iterable[Tuple[bytes, bytes]]) -> Dict[str, str]:
    out: Dict[str, str] = {}
    for k, v in headers:
        out[k.decode("latin-1").lower()] = v.decode("latin-1")
    return out


def _parse_query_string(qs: bytes) -> Dict[str, str]:
    if not qs:
        return {}
    from urllib.parse import parse_qs

    parsed = parse_qs(qs.decode("latin-1"), keep_blank_values=True)
    return {k: (v[0] if v else "") for k, v in parsed.items()}


def _client_ip_from_scope(scope: Mapping[str, Any]) -> str:
    client = scope.get("client")
    if isinstance(client, (list, tuple)) and client:
        return str(client[0])
    return ""


def _is_secure(scope: Mapping[str, Any], headers: Mapping[str, str]) -> bool:
    xf = headers.get("x-forwarded-proto", "")
    if "https" in xf:
        return True
    return str(scope.get("scheme") or "").lower() == "https"


def _absolute_url(scope: Mapping[str, Any], headers: Mapping[str, str], path: str) -> str:
    scheme = "https" if _is_secure(scope, headers) else "http"
    host = headers.get("host", "") or "localhost"
    return f"{scheme}://{host}{path}"


def _detect_client_ip(headers: Mapping[str, str], rules: Mapping[str, Any], remote_ip: str) -> str:
    ip = normalize_ip(remote_ip) or "0.0.0.0"

    cf_ranges = rules.get("cloudflare_cidrs")
    if isinstance(cf_ranges, list) and cf_ranges:
        for cidr in cf_ranges:
            if isinstance(cidr, str) and cidr and cidr_match(ip, cidr):
                cf_ip = normalize_ip(headers.get("cf-connecting-ip", ""))
                if cf_ip:
                    return cf_ip

                xff = headers.get("x-forwarded-for", "")
                if xff:
                    first = normalize_ip(xff.split(",")[0].strip())
                    if first:
                        return first
                break

    # fallback
    return ip


def _build_cookie(*, name: str, value: str, max_age: int, secure: bool) -> str:
    base = f"{name}={value}; Path=/; Max-Age={int(max_age)}; HttpOnly; SameSite=Lax"
    return base + ("; Secure" if secure else "")


async def _send_text(send: Callable[[Dict[str, Any]], Awaitable[None]], status: int, text: str):
    await send({
        "type": "http.response.start",
        "status": int(status),
        "headers": [(b"content-type", b"text/plain; charset=utf-8")],
    })
    await send({
        "type": "http.response.body",
        "body": text.encode("utf-8"),
        "more_body": False,
    })


async def _send_html(send: Callable[[Dict[str, Any]], Awaitable[None]], status: int, html: str):
    await send({
        "type": "http.response.start",
        "status": int(status),
        "headers": [(b"content-type", b"text/html; charset=utf-8")],
    })
    await send({
        "type": "http.response.body",
        "body": html.encode("utf-8"),
        "more_body": False,
    })


async def _send_redirect(send: Callable[[Dict[str, Any]], Awaitable[None]], status: int, location: str, *, extra_headers: Optional[List[Tuple[bytes, bytes]]] = None):
    headers: List[Tuple[bytes, bytes]] = [(b"location", location.encode("utf-8"))]
    if extra_headers:
        headers.extend(extra_headers)
    await send({
        "type": "http.response.start",
        "status": int(status),
        "headers": headers,
    })
    await send({
        "type": "http.response.body",
        "body": b"",
        "more_body": False,
    })


async def _respond_challenge(send: Callable[[Dict[str, Any]], Awaitable[None]], scope: Mapping[str, Any], headers: Mapping[str, str], license_key: str, api_url: str):
    base_url = api_url
    if base_url.lower().endswith("/api"):
        base_url = base_url[:-4]
    current_url = _absolute_url(scope, headers, str(scope.get("path") or "/"))
    from urllib.parse import urlencode

    # Preferred flow: opaque token (prevents license_key leak in URLs/logs)
    token = None
    try:
        import httpx

        async with httpx.AsyncClient(timeout=2.0) as client:
            r = await client.post(api_url.rstrip("/") + "/challenge/init", json={
                "license_key": license_key,
                "return_url": current_url,
            })
            if r.status_code >= 200 and r.status_code < 300:
                data = r.json()
                if isinstance(data, dict):
                    t = data.get("token")
                    if isinstance(t, str) and len(t) > 20:
                        token = t
    except Exception:
        token = None

    if token:
        location = base_url + "/security-check?" + urlencode({"token": token})
        await _send_redirect(send, 302, location)
        return

    location = base_url + "/security-check?" + urlencode({"license_key": license_key, "return_url": current_url})
    await _send_redirect(send, 302, location)


async def _respond_block(send: Callable[[Dict[str, Any]], Awaitable[None]], rules: Mapping[str, Any], ip: str, reason: str):
    cloak_html = rules.get("cloak_html")
    if isinstance(cloak_html, str) and cloak_html:
        await _send_html(send, 200, cloak_html)
        return

    rid = "RAY-" + os.urandom(6).hex()
    html = (
        "<!DOCTYPE html><html lang=\"en\"><head><meta charset=\"UTF-8\" /><meta name=\"viewport\" content=\"width=device-width, initial-scale=1.0\" />"
        "<title>Access Denied</title>"
        "<style>body{background:#050507;color:#d4d4d8;font-family:-apple-system,BlinkMacSystemFont,Segoe UI,Roboto,Helvetica,Arial,sans-serif;height:100vh;display:flex;flex-direction:column;justify-content:center;align-items:center;margin:0}.container{max-width:600px;padding:40px;text-align:center}h1{font-size:32px;font-weight:700;color:#fff;margin:0 0 10px}p{font-size:16px;line-height:1.5;color:#a1a1aa;margin-bottom:30px}.details{background:#0f0f11;border:1px solid #27272a;border-radius:8px;padding:15px;text-align:left;font-family:monospace;font-size:12px;color:#71717a;width:100%;box-sizing:border-box}.row{display:flex;justify-content:space-between;margin-bottom:5px}</style></head>"
        f"<body><div class=\"container\"><h1>Access Denied</h1><p>The owner of this website has banned your IP address ({ip}) based on global security policies.</p>"
        f"<div class=\"details\"><div class=\"row\"><span>Ray ID:</span><span>{rid}</span></div><div class=\"row\"><span>Your IP:</span><span>{ip}</span></div><div class=\"row\"><span>Reason:</span><span>{reason}</span></div></div></div></body></html>"
    )
    await _send_html(send, 403, html)
