import math
import threading
import time
import urllib.request
from http import cookiejar
from http.cookiejar import Cookie, CookieJar, CookiePolicy, DefaultCookiePolicy
from typing import TYPE_CHECKING

__all__ = ["IndexedCookieJar"]

# [RFC 1035](https://www.rfc-editor.org/rfc/rfc1035#section-2.3.4) caps DNS names at 255 octets;
# longer hosts fall back to scanning every stored domain.
MAX_HOST_LENGTH = 255


class IndexedCookieJar(CookieJar):
    """
    A `CookieJar` whose per-request cost does not grow with the number of stored cookies.

    Compared to `CookieJar`, on each request it:

    1. Looks up only the domains that can match the request host, instead of checking every stored domain.
    2. Skips the expiry sweep until the earliest known expiry has passed, instead of checking every stored cookie.
    """

    _cookies: dict[str, dict[str, dict[str, Cookie]]]
    _cookies_lock: threading.RLock
    _policy: CookiePolicy
    _next_expiry: float = math.inf

    if TYPE_CHECKING:

        def _cookies_for_domain(self, domain: str, request: urllib.request.Request) -> list[Cookie]: ...

    def set_cookie(self, cookie: Cookie) -> None:
        with self._cookies_lock:
            super().set_cookie(cookie)
            if cookie.expires is not None:
                self._next_expiry = min(self._next_expiry, cookie.expires)

    def clear_expired_cookies(self) -> None:
        if time.time() < self._next_expiry:
            return
        with self._cookies_lock:
            super().clear_expired_cookies()
            self._next_expiry = min((c.expires for c in self if c.expires is not None), default=math.inf)

    def _cookies_for_request(self, request: urllib.request.Request) -> list[Cookie]:
        hosts: tuple[str, str] = cookiejar.eff_request_host(request)  # type: ignore[attr-defined]
        uses_default_matching = type(self._policy).domain_return_ok is DefaultCookiePolicy.domain_return_ok
        if not uses_default_matching or max(map(len, hosts)) > MAX_HOST_LENGTH:
            return super()._cookies_for_request(request)  # type: ignore[misc, no-any-return]

        cookies: list[Cookie] = []
        for domain in _matching_domains(hosts):
            if domain in self._cookies:
                cookies.extend(self._cookies_for_domain(domain, request))
        return cookies


def _matching_domains(hosts: tuple[str, str]) -> dict[str, None]:
    """
    Every cookie domain `DefaultCookiePolicy.domain_return_ok` can accept for these hosts, least specific first.

    For `a.example.org` that is `""`, `.org`, `org`, `.example.org`, `example.org`, `.a.example.org`, `a.example.org`.
    """
    domains = {"": None}
    for host in hosts:
        dotted = host if host.startswith(".") else f".{host}"
        for index in range(len(dotted) - 1, -1, -1):
            if dotted[index] == ".":
                domains[dotted[index:]] = None
                domains[dotted[index + 1 :]] = None
    return domains