"""Optional JWKS-based JWT authentication for RoohAI endpoints.

Enabled by setting the ROOHAI_AUTH_JWKS_URL environment variable.
Requires the `pyjwt[crypto]` package: ``pip install roohai[auth]``.
"""

from __future__ import annotations

import logging
import os
from typing import Any

from fastapi import Depends, WebSocket
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer

logger = logging.getLogger(__name__)

_jwks_client = None
_bearer_scheme = HTTPBearer(auto_error=False)


def is_auth_enabled() -> bool:
    """Return True if JWKS-based auth is configured."""
    return bool(os.environ.get("ROOHAI_AUTH_JWKS_URL"))


def get_jwks_client():
    """Return a cached PyJWKClient, creating it on first call."""
    global _jwks_client
    if _jwks_client is not None:
        return _jwks_client

    try:
        from jwt import PyJWKClient
    except ImportError:
        raise ImportError(
            "ROOHAI_AUTH_JWKS_URL is set but PyJWT is not installed. "
            "Install it with: pip install roohai[auth]"
        )

    url = os.environ["ROOHAI_AUTH_JWKS_URL"]
    _jwks_client = PyJWKClient(url, cache_jwk_set=True, lifespan=3600)
    logger.info("JWKS client initialized: %s", url)
    return _jwks_client


def decode_token(token: str) -> dict[str, Any]:
    """Validate and decode a JWT using the JWKS endpoint.

    Raises jwt.exceptions.PyJWTError on any validation failure.
    """
    import jwt

    client = get_jwks_client()
    signing_key = client.get_signing_key_from_jwt(token)

    kwargs: dict[str, Any] = {
        "algorithms": ["RS256", "RS384", "RS512", "ES256", "ES384", "ES512"],
    }

    issuer = os.environ.get("ROOHAI_AUTH_ISSUER")
    if issuer:
        kwargs["issuer"] = issuer

    audience = os.environ.get("ROOHAI_AUTH_AUDIENCE")
    if audience:
        kwargs["audience"] = audience

    return jwt.decode(token, signing_key.key, **kwargs)


async def require_auth(
    credentials: HTTPAuthorizationCredentials | None = Depends(_bearer_scheme),
) -> dict[str, Any] | None:
    """FastAPI dependency: validate Bearer token if auth is enabled.

    Returns decoded claims on success, None if auth is disabled.
    Raises HTTP 401 if auth is enabled and the token is missing/invalid.
    """
    if not is_auth_enabled():
        return None

    from fastapi import HTTPException

    if not credentials:
        raise HTTPException(
            status_code=401,
            detail="Missing authentication token",
            headers={"WWW-Authenticate": "Bearer"},
        )

    try:
        return decode_token(credentials.credentials)
    except Exception as exc:
        logger.debug("JWT validation failed: %s", exc)
        raise HTTPException(
            status_code=401,
            detail="Invalid or expired token",
            headers={"WWW-Authenticate": "Bearer"},
        )


async def verify_ws_token(websocket: WebSocket) -> dict[str, Any] | None:
    """Validate JWT on a WebSocket connection.

    Checks the ``?token=`` query param first, then the ``Authorization`` header.
    Closes with code 4001 on failure. Returns claims or None (if auth disabled).
    """
    if not is_auth_enabled():
        return None

    token = websocket.query_params.get("token")
    if not token:
        auth_header = websocket.headers.get("authorization", "")
        if auth_header.lower().startswith("bearer "):
            token = auth_header[7:]

    if not token:
        await websocket.close(code=4001, reason="Missing authentication token")
        raise RuntimeError("WebSocket auth failed: no token")

    try:
        return decode_token(token)
    except Exception as exc:
        logger.debug("WebSocket JWT validation failed: %s", exc)
        await websocket.close(code=4001, reason="Invalid or expired token")
        raise RuntimeError(f"WebSocket auth failed: {exc}")
