from __future__ import annotations

import ipaddress
from datetime import datetime, timedelta
from typing import Set, Iterable

from pymongo import AsyncMongoClient, ReturnDocument
from micropie import HttpMiddleware


# ---------------------------------------------------------------------------
# IP helpers
# ---------------------------------------------------------------------------


def _valid_ip(value: str | None) -> str | None:
    """Parse and normalize an IP string, returning canonical string form or None."""
    try:
        if not value:
            return None
        return str(ipaddress.ip_address(value.strip()))
    except Exception:
        return None


# Cloudflare IPv4 + IPv6 ranges
# Source: https://www.cloudflare.com/ips/
_CLOUDFLARE_IP_RANGES: list[str] = [
    # IPv4
    "173.245.48.0/20",
    "103.21.244.0/22",
    "103.22.200.0/22",
    "103.31.4.0/22",
    "141.101.64.0/18",
    "108.162.192.0/18",
    "190.93.240.0/20",
    "188.114.96.0/20",
    "197.234.240.0/22",
    "198.41.128.0/17",
    "162.158.0.0/15",
    "104.16.0.0/13",
    "104.24.0.0/14",
    "172.64.0.0/13",
    "131.0.72.0/22",
    # IPv6
    "2400:cb00::/32",
    "2606:4700::/32",
    "2803:f800::/32",
    "2405:b500::/32",
    "2405:8100::/32",
    "2a06:98c0::/29",
    "2c0f:f248::/32",
]

_CLOUDFLARE_NETWORKS = [ipaddress.ip_network(n) for n in _CLOUDFLARE_IP_RANGES]


def _is_cloudflare_socket_ip(socket_ip: str | None) -> bool:
    """
    Returns True if the connecting socket IP belongs to Cloudflare's published ranges.
    This is the key anti-spoof check before trusting CF/XFF headers.
    """
    try:
        if not socket_ip:
            return False
        addr = ipaddress.ip_address(socket_ip)
        return any(addr in net for net in _CLOUDFLARE_NETWORKS)
    except Exception:
        return False


# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------


class MongoRateLimitMiddleware(HttpMiddleware):
    """
    Global MongoDB-based rate limiter with Cloudflare anti-spoofing.

    - One document per client key (IP, optionally IP+route bucket)
    - Fixed window counter
    - Escalating temporary blocks
    - Permanent block based on 24h violation history
    - Fully atomic (single DB op per request)
    - PyMongo Async API (no Motor)

    Security:
    - Only trusts CF/XFF headers if:
        1) trust_proxy_headers=True
        2) (optionally) CF-Ray is present when require_cf_ray=True
        3) the connecting socket IP (ASGI scope['client'][0]) is in Cloudflare IP ranges
      Otherwise, proxy headers are ignored.

    Notes:
    - Host allow-list is a sanity check, NOT sufficient to prevent origin bypass if the
      origin is publicly reachable.
    - Best defense: make origin only reachable from Cloudflare (Tunnel/firewall).
    """

    # --- rate config ---
    MAX_REQUESTS = 50
    WINDOW_SECONDS = 60

    BLOCK_AFTER_VIOLATIONS = 3
    BLOCK_FOR_SECONDS = 900

    PERMA_WINDOW_HOURS = 24
    PERMA_BLOCK_AFTER = 10

    def __init__(
        self,
        mongo_uri: str,
        db_name: str,
        collection_name: str = "rate_limits_global",
        *,
        allowed_hosts: Set[str] | None = None,
        trust_proxy_headers: bool = True,
        require_cf_ray: bool = True,
        limit_methods: Iterable[str] | None = None,
        # Optional: include METHOD+PATH in key (lets you set different limits by route)
        bucket_by_route: bool = False,
        # If we can't reliably identify the client, return 403 (recommended)
        fail_closed: bool = True,
    ):
        self.client = AsyncMongoClient(mongo_uri)
        self.db = self.client[db_name]
        self.collection = self.db[collection_name]

        self.allowed_hosts = set(h.lower() for h in (allowed_hosts or set()))
        self.trust_proxy_headers = trust_proxy_headers
        self.require_cf_ray = require_cf_ray
        self.limit_methods = (
            {method.upper() for method in limit_methods} if limit_methods else None
        )

        self.bucket_by_route = bucket_by_route
        self.fail_closed = fail_closed

    # ---------------------------------------------------------
    # Client identification
    # ---------------------------------------------------------

    def _host_allowed(self, headers: dict) -> bool:
        if not self.allowed_hosts:
            return True
        host = (headers.get("host") or "").split(":", 1)[0].lower()
        return bool(host) and host in self.allowed_hosts

    def _socket_ip(self, request) -> str | None:
        client = (request.scope or {}).get("client") or (None, 0)
        return _valid_ip(client[0])

    def _can_trust_proxy_headers(self, headers: dict, socket_ip: str | None) -> bool:
        if not self.trust_proxy_headers:
            return False
        if self.require_cf_ray and not headers.get("cf-ray"):
            return False
        # Critical anti-spoof: only trust if the peer socket IP is Cloudflare
        return _is_cloudflare_socket_ip(socket_ip)

    def _client_ip(self, request) -> str | None:
        headers = getattr(request, "headers", {}) or {}

        # Optional host sanity check
        if not self._host_allowed(headers):
            return None

        socket_ip = self._socket_ip(request)
        can_trust = self._can_trust_proxy_headers(headers, socket_ip)

        if can_trust:
            # Cloudflare headers (preferred)
            for h in ("cf-connecting-ip", "true-client-ip"):
                ip = _valid_ip(headers.get(h))
                if ip:
                    return ip

            # Standard proxy chain
            xff = headers.get("x-forwarded-for")
            if isinstance(xff, str):
                ip = _valid_ip(xff.split(",", 1)[0])
                if ip:
                    return ip

            # Fallback proxy header
            ip = _valid_ip(headers.get("x-real-ip"))
            if ip:
                return ip

        # Untrusted path: fall back to socket peer IP (Heroku router, etc.)
        return socket_ip

    def _key(self, client_ip: str, request) -> str:
        if not self.bucket_by_route:
            return client_ip
        method = (
            getattr(request, "method", None)
            or (request.scope or {}).get("method")
            or "GET"
        ).upper()
        path = (request.scope or {}).get("path") or "/"
        return f"{client_ip}|{method}|{path}"

    # ---------------------------------------------------------
    # Middleware hook
    # ---------------------------------------------------------

    async def before_request(self, request):
        method = (
            getattr(request, "method", None)
            or (request.scope or {}).get("method")
            or "GET"
        ).upper()
        if self.limit_methods is not None and method not in self.limit_methods:
            return None

        client_ip = self._client_ip(request)

        # Avoid collapsing unknowns into a shared key (DoS vector)
        if not client_ip:
            if self.fail_closed:
                return {"status_code": 403, "body": "Forbidden.", "headers": []}
            return None  # fail open (not recommended)

        now = datetime.utcnow()

        window_start_cutoff = now - timedelta(seconds=self.WINDOW_SECONDS)
        perma_window_cutoff = now - timedelta(hours=self.PERMA_WINDOW_HOURS)

        key = self._key(client_ip, request)

        doc = await self.collection.find_one_and_update(
            {"_id": key},
            [
                # 1) Baseline fields
                {
                    "$set": {
                        "_id": key,
                        "ip": client_ip,
                        "count": {"$ifNull": ["$count", 0]},
                        "window_start": {"$ifNull": ["$window_start", now]},
                        "violations": {"$ifNull": ["$violations", 0]},
                        "blocked_until": {"$ifNull": ["$blocked_until", None]},
                        "permanent_blocked": {"$ifNull": ["$permanent_blocked", False]},
                        "permanent_blocked_at": {
                            "$ifNull": ["$permanent_blocked_at", None]
                        },
                        "violation_events": {"$ifNull": ["$violation_events", []]},
                    }
                },
                # 2) Prune old violation events
                {
                    "$set": {
                        "violation_events": {
                            "$filter": {
                                "input": "$violation_events",
                                "as": "t",
                                "cond": {"$gte": ["$$t", perma_window_cutoff]},
                            }
                        }
                    }
                },
                # 3) Are we currently blocked?
                {
                    "$set": {
                        "_blocked_now": {
                            "$or": [
                                "$permanent_blocked",
                                {
                                    "$and": [
                                        {"$ne": ["$blocked_until", None]},
                                        {"$gt": ["$blocked_until", now]},
                                    ]
                                },
                            ]
                        }
                    }
                },
                # 4) Update window/count (only if not blocked)
                {
                    "$set": {
                        "_window_expired": {
                            "$cond": [
                                "$_blocked_now",
                                False,
                                {"$lt": ["$window_start", window_start_cutoff]},
                            ]
                        }
                    }
                },
                {
                    "$set": {
                        "window_start": {
                            "$cond": [
                                "$_blocked_now",
                                "$window_start",
                                {"$cond": ["$_window_expired", now, "$window_start"]},
                            ]
                        },
                        "count": {
                            "$cond": [
                                "$_blocked_now",
                                "$count",
                                {
                                    "$cond": [
                                        "$_window_expired",
                                        1,
                                        {"$add": ["$count", 1]},
                                    ]
                                },
                            ]
                        },
                    }
                },
                # 5) Over limit?
                {
                    "$set": {
                        "_over_limit": {
                            "$and": [
                                {"$not": "$_blocked_now"},
                                {"$gt": ["$count", self.MAX_REQUESTS]},
                            ]
                        }
                    }
                },
                # 6) Record violation if over limit
                {
                    "$set": {
                        "violations": {
                            "$cond": [
                                "$_over_limit",
                                {"$add": ["$violations", 1]},
                                "$violations",
                            ]
                        },
                        "violation_events": {
                            "$cond": [
                                "$_over_limit",
                                {"$concatArrays": ["$violation_events", [now]]},
                                "$violation_events",
                            ]
                        },
                    }
                },
                # 7) Temporary block escalation
                {
                    "$set": {
                        "blocked_until": {
                            "$cond": [
                                {
                                    "$and": [
                                        "$_over_limit",
                                        {
                                            "$gte": [
                                                "$violations",
                                                self.BLOCK_AFTER_VIOLATIONS,
                                            ]
                                        },
                                    ]
                                },
                                now + timedelta(seconds=self.BLOCK_FOR_SECONDS),
                                "$blocked_until",
                            ]
                        }
                    }
                },
                # 8) Permanent block escalation
                {"$set": {"_events_24h": {"$size": "$violation_events"}}},
                {
                    "$set": {
                        "permanent_blocked": {
                            "$cond": [
                                {
                                    "$and": [
                                        "$_over_limit",
                                        {
                                            "$gte": [
                                                "$_events_24h",
                                                self.PERMA_BLOCK_AFTER,
                                            ]
                                        },
                                    ]
                                },
                                True,
                                "$permanent_blocked",
                            ]
                        },
                        "permanent_blocked_at": {
                            "$cond": [
                                {
                                    "$and": [
                                        "$_over_limit",
                                        {
                                            "$gte": [
                                                "$_events_24h",
                                                self.PERMA_BLOCK_AFTER,
                                            ]
                                        },
                                        {"$eq": ["$permanent_blocked_at", None]},
                                    ]
                                },
                                now,
                                "$permanent_blocked_at",
                            ]
                        },
                    }
                },
                # 9) Cleanup temp fields
                {
                    "$unset": [
                        "_blocked_now",
                        "_window_expired",
                        "_over_limit",
                        "_events_24h",
                    ]
                },
            ],
            upsert=True,
            return_document=ReturnDocument.AFTER,
            projection={"count": 1, "blocked_until": 1, "permanent_blocked": 1},
        )

        doc = doc or {}

        # --- responses ---
        if doc.get("permanent_blocked"):
            return {
                "status_code": 403,
                "body": f"Access permanently blocked for IP {client_ip}.",
                "headers": [],
            }

        blocked_until = doc.get("blocked_until")
        if isinstance(blocked_until, datetime) and now < blocked_until:
            retry_after = max(0, int((blocked_until - now).total_seconds()))
            return {
                "status_code": 429,
                "body": f"Too many requests from {client_ip}. Temporarily blocked.",
                "headers": [("Retry-After", str(retry_after))],
            }

        if int(doc.get("count", 0)) > self.MAX_REQUESTS:
            return {
                "status_code": 429,
                "body": f"Rate limit exceeded for IP {client_ip}.",
                "headers": [],
            }

        return None

    async def after_request(self, request, status_code, response_body, extra_headers):
        return None
