from functools import lru_cache
from typing import Any, Dict, Optional
from urllib.parse import urljoin

import httpx
from fastapi import HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from jose import JWTError, jwk, jwt
from pydantic import BaseModel, Field, HttpUrl
from pydantic_settings import BaseSettings

security = HTTPBearer(auto_error=False)

class Settings(BaseSettings):
    # OIDC issuer (discovery URL root)
    oidc_issuer: HttpUrl = Field(..., env="OIDC_ISSUER")
    # Optional audience: if set the backend will validate the token 'aud' claim.
    # If omitted, audience validation is skipped (use with caution).
    oidc_audience: Optional[str] = Field(None, env="OIDC_AUDIENCE")
    oidc_client_id: str = Field(..., env="OIDC_CLIENT_ID")
    oidc_client_secret: Optional[str] = Field(None, env="OIDC_CLIENT_SECRET")
    oidc_verify_iss: bool = Field(True, env="OIDC_VERIFY_ISS")

    class Config:
        env_file = ".env"

settings = Settings()

class TokenClaims(BaseModel):
    iss: str
    sub: str
    aud: Any
    exp: int
    iat: Optional[int]
    azp: Optional[str]
    email: Optional[str]
    preferred_username: Optional[str]

@lru_cache(maxsize=1)
def get_discovery() -> Dict[str, Any]:
    issuer = str(settings.oidc_issuer).rstrip("/")
    discovery_url = urljoin(issuer + "/", ".well-known/openid-configuration")
    with httpx.Client(timeout=10.0) as client:
        response = client.get(discovery_url)
        if response.status_code != 200:
            raise RuntimeError(f"Unable to load OIDC discovery document from {discovery_url}")
        return response.json()

@lru_cache(maxsize=1)
def get_jwks() -> Dict[str, Any]:
    discovery = get_discovery()
    jwks_uri = discovery.get("jwks_uri")
    if not jwks_uri:
        raise RuntimeError("OIDC discovery document does not contain jwks_uri")
    with httpx.Client(timeout=10.0) as client:
        response = client.get(jwks_uri)
        if response.status_code != 200:
            raise RuntimeError(f"Unable to load JWKS from {jwks_uri}")
        return response.json()

def get_signing_key(token: str) -> Dict[str, Any]:
    try:
        headers = jwt.get_unverified_header(token)
    except JWTError as exc:
        raise HTTPException(status_code=401, detail="Invalid JWT header") from exc

    kid = headers.get("kid")
    if not kid:
        raise HTTPException(status_code=401, detail="JWT missing kid header")

    jwks = get_jwks().get("keys", [])
    for key in jwks:
        if key.get("kid") == kid:
            return key

    raise HTTPException(status_code=401, detail="Unable to find matching JWKS key")

async def validate_token(credentials: HTTPAuthorizationCredentials = Security(security)) -> TokenClaims:
    if not credentials or credentials.scheme.lower() != "bearer":
        raise HTTPException(status_code=401, detail="Missing Bearer authorization header")

    token = credentials.credentials
    key = get_signing_key(token)
    public_key = jwk.construct(key)
    try:
        # Only pass `audience` to the decoder if configured. Some providers
        # (or local development setups) may not include the aud claim in a
        # way that matches your API identifier; in that case leave
        # `OIDC_AUDIENCE` unset and the audience check will be skipped.
        jwt_kwargs = {
            "algorithms": [key.get("alg", "RS256")],
            "issuer": discovery.get("issuer") if settings.oidc_verify_iss else None,
        }
        if settings.oidc_audience:
            jwt_kwargs["audience"] = settings.oidc_audience

        verified = jwt.decode(token, public_key, **jwt_kwargs)
    except JWTError as exc:
        raise HTTPException(status_code=401, detail="Invalid or expired token") from exc

    try:
        return TokenClaims(**verified)
    except Exception as exc:
        raise HTTPException(status_code=401, detail="Unable to parse token claims") from exc