from __future__ import annotations

import hashlib
import json
import os
import threading
import time
from dataclasses import dataclass
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.request import Request, urlopen

@dataclass(frozen=True)
class Auth0Settings:
    domain: str
    allowed_email_domains: frozenset[str]
    allowed_connections: frozenset[str]
    userinfo_timeout_seconds: float = 5.0
    userinfo_cache_seconds: float = 300.0

    @property
    def issuer(self) -> str:
        return f"https://{self.domain}/"

    @property
    def userinfo_url(self) -> str:
        return f"{self.issuer}userinfo"

    @classmethod
    def from_env(cls) -> Auth0Settings:
        domain = os.environ.get("AUTH0_DOMAIN", "").strip().removeprefix("https://").rstrip("/")
        if not domain:
            raise RuntimeError("Missing Auth0 configuration: AUTH0_DOMAIN")

        allowed_domains = frozenset(
            value.strip().lower().lstrip("@")
            for value in os.environ.get("AUTH0_ALLOWED_EMAIL_DOMAINS", "pingcap.com").split(",")
            if value.strip()
        )
        allowed_connections = frozenset(
            value.strip()
            for value in os.environ.get("AUTH0_ALLOWED_CONNECTIONS", "google-oauth2").split(",")
            if value.strip()
        )
        if not allowed_domains or not allowed_connections:
            raise RuntimeError("Auth0 email-domain and connection allowlists must not be empty")

        return cls(
            domain=domain,
            allowed_email_domains=allowed_domains,
            allowed_connections=allowed_connections,
        )


@dataclass(frozen=True)
class AuthPrincipal:
    subject: str
    email: str
    email_verified: bool
    name: str
    picture_url: str | None
    connection: str


class AuthError(Exception):
    def __init__(self, status: int, code: str, message: str) -> None:
        super().__init__(message)
        self.status = status
        self.code = code
        self.message = message


class Auth0Authenticator:
    def __init__(self, settings: Auth0Settings) -> None:
        self.settings = settings
        self._userinfo_cache: dict[str, tuple[float, dict[str, Any]]] = {}
        self._cache_lock = threading.Lock()

    def authenticate(self, authorization_header: str | None) -> AuthPrincipal:
        token = self._bearer_token(authorization_header)
        profile = self._userinfo(token)

        subject = str(profile.get("sub") or "").strip()
        if not subject:
            raise AuthError(401, "invalid_access_token", "The token subject is invalid")

        connection = subject.partition("|")[0]
        if connection not in self.settings.allowed_connections:
            raise AuthError(403, "account_not_allowed", "Only the configured Google connection is allowed")

        email = str(profile.get("email") or "").strip().lower()
        email_domain = email.rpartition("@")[2]
        if not email or email_domain not in self.settings.allowed_email_domains:
            raise AuthError(403, "account_not_allowed", "A PingCAP Google account is required")
        if profile.get("email_verified") is not True:
            raise AuthError(403, "account_not_allowed", "A verified email address is required")

        name = str(profile.get("name") or profile.get("nickname") or email.split("@", 1)[0]).strip()
        picture_url = str(profile.get("picture") or "").strip() or None
        return AuthPrincipal(
            subject=subject,
            email=email,
            email_verified=True,
            name=name,
            picture_url=picture_url,
            connection=connection,
        )

    @staticmethod
    def _bearer_token(authorization_header: str | None) -> str:
        scheme, separator, token = (authorization_header or "").strip().partition(" ")
        if not separator or scheme.lower() != "bearer" or not token.strip():
            raise AuthError(401, "authentication_required", "A Bearer access token is required")
        return token.strip()

    def _userinfo(self, token: str) -> dict[str, Any]:
        cache_key = hashlib.sha256(token.encode("utf-8")).hexdigest()
        now = time.monotonic()
        with self._cache_lock:
            cached = self._userinfo_cache.get(cache_key)
            if cached and cached[0] > now:
                return cached[1]

        request = Request(
            self.settings.userinfo_url,
            headers={"Authorization": f"Bearer {token}", "Accept": "application/json"},
        )
        try:
            with urlopen(request, timeout=self.settings.userinfo_timeout_seconds) as response:
                profile = json.loads(response.read().decode("utf-8"))
        except HTTPError as exc:
            if exc.code in {401, 403}:
                raise AuthError(401, "invalid_access_token", "Auth0 rejected the access token") from exc
            raise AuthError(503, "auth_provider_unavailable", "Auth0 user information is unavailable") from exc
        except (URLError, TimeoutError, json.JSONDecodeError) as exc:
            raise AuthError(503, "auth_provider_unavailable", "Auth0 user information is unavailable") from exc
        if not isinstance(profile, dict):
            raise AuthError(503, "auth_provider_unavailable", "Auth0 returned an invalid user profile")

        with self._cache_lock:
            self._userinfo_cache[cache_key] = (
                now + self.settings.userinfo_cache_seconds,
                profile,
            )
            if len(self._userinfo_cache) > 512:
                self._userinfo_cache = {
                    key: value for key, value in self._userinfo_cache.items() if value[0] > now
                }
        return profile