"""OIDC authentication configuration and helpers.

Provides :class:`OIDCConfig` for validated OIDC configuration from
environment variables, PKCE helpers for the authorization code flow,
and session cookie minting/validation utilities.

See ``designs/OIDC_AUTH.md`` for the full design.
"""

from __future__ import annotations

import base64
import hashlib
import hmac
import logging
import os
import secrets
import time
from dataclasses import dataclass

import httpx
import jwt

_logger = logging.getLogger(__name__)

# ── PKCE helpers (RFC 7636) ──────────────────────────────────────


def generate_code_verifier() -> str:
    """Generate a PKCE code verifier (43–128 URL-safe chars).

    :returns: A random code verifier string suitable for the
        ``code_verifier`` parameter in the OIDC token exchange.
    """
    return secrets.token_urlsafe(64)[:96]


def derive_code_challenge(code_verifier: str) -> str:
    """Derive a PKCE S256 code challenge from a code verifier.

    :param code_verifier: The code verifier string generated by
        :func:`generate_code_verifier`.
    :returns: Base64url-encoded SHA256 hash of the verifier, with
        padding stripped (per RFC 7636 Appendix B).
    """
    digest = hashlib.sha256(code_verifier.encode("ascii")).digest()
    return base64.urlsafe_b64encode(digest).rstrip(b"=").decode("ascii")


# ── Session cookie helpers ───────────────────────────────────────


def mint_session_token(
    user_id: str,
    cookie_secret: bytes,
    ttl_seconds: int,
    provider: str,
) -> str:
    """
    Mint a signed session JWT with a second-granularity lifetime.

    The seconds-based core behind :func:`mint_session_cookie`. A managed
    runner needs a short-lived (sub-hour) owner token, which the hours-only
    cookie helper cannot express; both share this HS256 claim shape so the
    same validator (:meth:`UnifiedAuthProvider._check_cookie`) accepts
    either.

    :param user_id: The authenticated user's email, e.g.
        ``"alice@example.com"``.
    :param cookie_secret: HMAC key for HS256 signing.
    :param ttl_seconds: Token lifetime in seconds.
    :param provider: Identity provider name, e.g. ``"google"`` or
        ``"accounts"``. Stored as an informational claim.
    :returns: An HS256-signed JWT string.
    """
    now = int(time.time())
    payload = {
        "sub": user_id,
        "iat": now,
        "exp": now + ttl_seconds,
        "provider": provider,
    }
    return jwt.encode(payload, cookie_secret, algorithm="HS256")


def mint_session_cookie(
    user_id: str,
    cookie_secret: bytes,
    ttl_hours: int,
    provider: str,
) -> str:
    """Mint a signed session cookie JWT.

    :param user_id: The authenticated user's email, e.g.
        ``"alice@example.com"``.
    :param cookie_secret: HMAC key for HS256 signing.
    :param ttl_hours: Session lifetime in hours.
    :param provider: Identity provider name, e.g. ``"google"``
        or ``"github"``. Stored as an informational claim.
    :returns: An HS256-signed JWT string.
    """
    return mint_session_token(user_id, cookie_secret, ttl_hours * 3600, provider)


def hmac_digest(token: str, secret: bytes) -> str:
    """Compute an HMAC-SHA256 digest of a cookie token for cache keying.

    :param token: The raw cookie JWT string.
    :param secret: The cookie secret used as the HMAC key.
    :returns: Hex-encoded HMAC-SHA256 digest.
    """
    return hmac.new(secret, token.encode("utf-8"), hashlib.sha256).hexdigest()


# ── GitHub endpoint constants ────────────────────────────────────


_GITHUB_ISSUER = "https://github.com"
_GITHUB_AUTHORIZATION_ENDPOINT = "https://github.com/login/oauth/authorize"
_GITHUB_TOKEN_ENDPOINT = "https://github.com/login/oauth/access_token"
_GITHUB_USERINFO_ENDPOINT = "https://api.github.com/user"
_GITHUB_EMAILS_ENDPOINT = "https://api.github.com/user/emails"
_GITHUB_SCOPES = "read:user user:email"


# ── OIDCConfig ───────────────────────────────────────────────────


@dataclass(frozen=True)
class OIDCConfig:
    """Validated OIDC configuration, constructed once at startup.

    All fields are validated at construction time via
    :meth:`from_env`. Missing or invalid values cause immediate
    startup failure.

    :param issuer: OIDC issuer URL, e.g.
        ``"https://accounts.google.com"``.
    :param client_id: OAuth client ID registered with the IdP.
    :param client_secret: OAuth client secret.
    :param redirect_uri: Full callback URL, e.g.
        ``"https://myapp.example.com/auth/callback"``.
    :param cookie_secret: HMAC key for session cookie signing
        (at least 32 bytes).
    :param scopes: Space-separated OAuth scopes, e.g.
        ``"openid email profile"``.
    :param session_ttl_hours: Session cookie lifetime in hours.
    :param logout_redirect_uri: Optional IdP end-session endpoint
        URL. ``None`` means logout just clears the cookie.
    :param allowed_domains: Optional frozenset of allowed email
        domains, e.g. ``frozenset({"example.com"})``. ``None``
        means all domains are accepted.
    :param provider_type: ``"oidc"`` for standard OIDC or
        ``"github"`` for GitHub OAuth.
    :param authorization_endpoint: IdP authorization URL.
    :param token_endpoint: IdP token exchange URL.
    :param jwks_uri: IdP JWKS URL for ``id_token`` validation.
        ``None`` for GitHub (no ``id_token``).
    :param userinfo_endpoint: Userinfo URL for providers that
        don't issue ``id_token`` (GitHub). ``None`` for standard
        OIDC providers.
    :param skip_email_verification: When ``True``, accept the
        ``id_token`` email claim without requiring
        ``email_verified``. Only affects the generic-OIDC path;
        GitHub always requires a verified primary email.
    """

    issuer: str
    client_id: str
    client_secret: str
    redirect_uri: str
    cookie_secret: bytes
    scopes: str
    session_ttl_hours: int
    logout_redirect_uri: str | None
    allowed_domains: frozenset[str] | None
    provider_type: str
    authorization_endpoint: str
    token_endpoint: str
    jwks_uri: str | None
    userinfo_endpoint: str | None
    allow_invites: bool
    skip_email_verification: bool = False

    @property
    def base_url(self) -> str:
        """Public base URL of the deployment, derived from ``redirect_uri``.

        The redirect URI is the full callback
        (``https://app.example.com/auth/callback``); the base URL is its
        scheme + host (+ port). Used to build OIDC invite links
        (``<base_url>/auth/login?invite=...``).

        :returns: The scheme://host[:port] prefix, e.g.
            ``"https://app.example.com"``.
        """
        from urllib.parse import urlsplit, urlunsplit

        parts = urlsplit(self.redirect_uri)
        return urlunsplit((parts.scheme, parts.netloc, "", "", ""))

    @property
    def secure_cookies(self) -> bool:
        """Whether to set ``Secure`` flag and use ``__Host-`` prefix.

        Derived from the redirect URI scheme: ``True`` for HTTPS,
        ``False`` for HTTP (local development). The ``__Host-``
        cookie prefix requires HTTPS — browsers silently drop the
        cookie on plain HTTP, causing an infinite login redirect.

        :returns: ``True`` when cookies should be secure.
        """
        return self.redirect_uri.startswith("https://")

    @property
    def session_cookie_name(self) -> str:
        """Cookie name for the session JWT.

        Uses the ``__Host-`` prefix on HTTPS (prevents cookie
        tossing attacks) and a plain name on HTTP (local dev).

        :returns: Cookie name string.
        """
        return "__Host-ap_session" if self.secure_cookies else "ap_session"

    @staticmethod
    def from_env() -> OIDCConfig:
        """Read and validate all OIDC env vars.

        Fetches the OIDC discovery document for standard providers,
        or uses hardcoded endpoints for GitHub.

        :returns: A validated :class:`OIDCConfig`.
        :raises RuntimeError: If any required env var is missing or
            invalid, or if OIDC discovery fails.
        """

        def _require(name: str) -> str:
            val = os.environ.get(name, "").strip()
            if not val:
                raise RuntimeError(
                    f"Missing required environment variable {name} "
                    f"(OMNIGENT_AUTH_PROVIDER=oidc requires it)"
                )
            return val

        issuer = _require("OMNIGENT_OIDC_ISSUER")
        client_id = _require("OMNIGENT_OIDC_CLIENT_ID")
        client_secret = _require("OMNIGENT_OIDC_CLIENT_SECRET")
        # Redirect URI: an explicit value wins; otherwise derive it from
        # OMNIGENT_DOMAIN (the same var the Caddy HTTPS overlay uses) as
        # ``https://<domain>/auth/callback``. A domain-based deploy then
        # sets one var instead of two, and can't get the http/https scheme
        # wrong. Must still match the IdP app's callback exactly.
        redirect_uri = os.environ.get("OMNIGENT_OIDC_REDIRECT_URI", "").strip()
        if not redirect_uri:
            domain = os.environ.get("OMNIGENT_DOMAIN", "").strip()
            if not domain:
                raise RuntimeError(
                    "Missing required environment variable OMNIGENT_OIDC_REDIRECT_URI "
                    "(or set OMNIGENT_DOMAIN to derive https://<domain>/auth/callback)"
                )
            redirect_uri = f"https://{domain}/auth/callback"

        cookie_secret_hex = _require("OMNIGENT_OIDC_COOKIE_SECRET")
        try:
            cookie_secret = bytes.fromhex(cookie_secret_hex)
        except ValueError as exc:
            raise RuntimeError("OMNIGENT_OIDC_COOKIE_SECRET must be a valid hex string") from exc
        if len(cookie_secret) < 32:
            raise RuntimeError(
                "OMNIGENT_OIDC_COOKIE_SECRET must be at least 32 bytes (64 hex chars)"
            )

        session_ttl_hours = int(os.environ.get("OMNIGENT_OIDC_SESSION_TTL_HOURS", "8"))
        logout_redirect_uri = (
            os.environ.get("OMNIGENT_OIDC_LOGOUT_REDIRECT_URI", "").strip() or None
        )

        raw_domains = os.environ.get("OMNIGENT_OIDC_ALLOWED_DOMAINS", "").strip()
        allowed_domains: frozenset[str] | None = None
        if raw_domains:
            allowed_domains = frozenset(
                d.strip().lower() for d in raw_domains.split(",") if d.strip()
            )

        # Opt-in individual invites (admin pre-authorizes an off-domain
        # email; see omnigent/server/oidc_access.py). Off by default.
        # Imported lazily to avoid a circular import (auth imports oidc).
        from omnigent.server.auth import env_var_is_truthy

        allow_invites = env_var_is_truthy("OMNIGENT_OIDC_ALLOW_INVITES")

        # Some IdPs (e.g. Okta without custom API Access Management)
        # omit ``email_verified`` for directory-provisioned users even
        # though the directory is authoritative for the address. This
        # opt-out trusts any signed email claim from the IdP — only
        # enable it when the issuer is a trusted enterprise directory.
        skip_email_verification = env_var_is_truthy("OMNIGENT_OIDC_SKIP_EMAIL_VERIFICATION")
        if skip_email_verification:
            _logger.warning(
                "OMNIGENT_OIDC_SKIP_EMAIL_VERIFICATION is set: the "
                "email_verified claim will not be required on OIDC "
                "id_tokens. Any signed email claim from %s will be "
                "trusted as the user's identity.",
                issuer,
            )

        # Determine provider type and resolve endpoints.
        is_github = issuer.rstrip("/") == _GITHUB_ISSUER

        if is_github:
            # Empty string (forwarded by `${VAR:-}` wrappers) → default.
            scopes = (os.environ.get("OMNIGENT_OIDC_SCOPES") or _GITHUB_SCOPES).strip()
            return OIDCConfig(
                issuer=_GITHUB_ISSUER,
                client_id=client_id,
                client_secret=client_secret,
                redirect_uri=redirect_uri,
                cookie_secret=cookie_secret,
                scopes=scopes,
                session_ttl_hours=session_ttl_hours,
                logout_redirect_uri=logout_redirect_uri,
                allowed_domains=allowed_domains,
                provider_type="github",
                authorization_endpoint=_GITHUB_AUTHORIZATION_ENDPOINT,
                token_endpoint=_GITHUB_TOKEN_ENDPOINT,
                jwks_uri=None,
                userinfo_endpoint=_GITHUB_USERINFO_ENDPOINT,
                allow_invites=allow_invites,
                skip_email_verification=skip_email_verification,
            )

        # Standard OIDC: fetch discovery document.
        scopes = (os.environ.get("OMNIGENT_OIDC_SCOPES") or "openid email profile").strip()
        discovery_url = issuer.rstrip("/") + "/.well-known/openid-configuration"
        try:
            resp = httpx.get(discovery_url, timeout=10.0)
            resp.raise_for_status()
            doc = resp.json()
        except Exception as exc:
            raise RuntimeError(
                f"Failed to fetch OIDC discovery document from {discovery_url}: {exc}"
            ) from exc

        authorization_endpoint = doc.get("authorization_endpoint")
        token_endpoint = doc.get("token_endpoint")
        jwks_uri = doc.get("jwks_uri")

        if not authorization_endpoint or not token_endpoint or not jwks_uri:
            raise RuntimeError(
                f"OIDC discovery document at {discovery_url} missing "
                f"required fields (authorization_endpoint, "
                f"token_endpoint, jwks_uri)"
            )

        return OIDCConfig(
            issuer=issuer,
            client_id=client_id,
            client_secret=client_secret,
            redirect_uri=redirect_uri,
            cookie_secret=cookie_secret,
            scopes=scopes,
            session_ttl_hours=session_ttl_hours,
            logout_redirect_uri=logout_redirect_uri,
            allowed_domains=allowed_domains,
            provider_type="oidc",
            authorization_endpoint=authorization_endpoint,
            token_endpoint=token_endpoint,
            jwks_uri=jwks_uri,
            userinfo_endpoint=doc.get("userinfo_endpoint"),
            allow_invites=allow_invites,
            skip_email_verification=skip_email_verification,
        )