"""Keycloak JWT authentication middleware."""
from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Optional

import httpx
from dotenv import load_dotenv
from fastapi import Depends, HTTPException, Request, status
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwt
from jose.constants import Algorithms

load_dotenv(Path(__file__).parent.parent / ".env", override=True)

logger = logging.getLogger(__name__)
security = HTTPBearer(auto_error=False)

KEYCLOAK_SERVER = os.getenv("KEYCLOAK_SERVER_URL", "https://keycloak.plfai.cn")
KEYCLOAK_REALM = os.getenv("KEYCLOAK_REALM", "plfai")
KEYCLOAK_AUDIENCE = os.getenv("KEYCLOAK_AUDIENCE", "account")

OIDC_CONFIG_URL = f"{KEYCLOAK_SERVER}/realms/{KEYCLOAK_REALM}/.well-known/openid-configuration"

# Cached OIDC config + JWKS
_oidc_config: Optional[dict] = None
_jwks: Optional[dict] = None


async def _fetch_oidc_config() -> tuple[dict, dict]:
    """Fetch OIDC config and JWKS, caching them."""
    global _oidc_config, _jwks

    if _oidc_config is None:
        async with httpx.AsyncClient() as client:
            resp = await client.get(OIDC_CONFIG_URL)
            resp.raise_for_status()
            _oidc_config = resp.json()

            jwks_resp = await client.get(_oidc_config["jwks_uri"])
            jwks_resp.raise_for_status()
            _jwks = jwks_resp.json()

    return _oidc_config, _jwks  # type: ignore[return-value]


async def get_current_user(
    request: Request,
    credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
) -> dict:
    """Validate Bearer token and return user claims. Raises 401 if invalid."""
    if credentials is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing Authorization header",
        )

    token = credentials.credentials

    try:
        _, jwks = await _fetch_oidc_config()
    except Exception as e:
        logger.error(f"Failed to fetch OIDC config: {e}")
        raise HTTPException(
            status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
            detail="Auth service unavailable",
        )

    # Get kid from token header
    try:
        unverified = jwt.get_unverified_header(token)
    except JWTError:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Invalid token header",
        )

    kid = unverified.get("kid")
    if not kid:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Missing kid in token",
        )

    # Find matching public key
    key = None
    for k in jwks.get("keys", []):
        if k.get("kid") == kid:
            key = k
            break

    if key is None:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Unknown signing key",
        )

    # Verify signature + expiry
    try:
        payload = jwt.decode(
            token,
            key,
            algorithms=[Algorithms.RS256],
            audience=KEYCLOAK_AUDIENCE,
            options={"verify_exp": True},
        )
        return payload
    except JWTError as e:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail=f"Invalid token: {e}",
        )