import jwt
from typing import List
from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.security import OAuth2AuthorizationCodeBearer, SecurityScopes
from jwt import PyJWKClient

TENANT_ID = "your-tenant-id"
API_CLIENT_ID = "your-backend-api-client-id" # The Application ID of the API itself
JWKS_URL = f"https://login.microsoftonline.com/{TENANT_ID}/discovery/v2.0/keys"

jwks_client = PyJWKClient(JWKS_URL)

# Native security scheme enables the top-right "Authorize" button in Swagger UI
oauth2_scheme = OAuth2AuthorizationCodeBearer(
    authorizationUrl=f"https://login.microsoftonline.com/{TENANT_ID}/oauth2/v2.0/authorize",
    tokenUrl=f"https://login.microsoftonline.com/{TENANT_ID}/oauth2/v2.0/token",
)

class EntraRoleChecker:
    def __init__(self, required_roles: List[str]):
        self.required_roles = required_roles

    def __call__(self, token: str = Depends(oauth2_scheme)):
        try:
            signing_key = jwks_client.get_signing_key_from_jwt(token)
            payload = jwt.decode(
                token,
                signing_key.key,
                algorithms=["RS256"],
                audience=API_CLIENT_ID, # Ensures token was meant for THIS API
                issuer=f"https://sts.windows.net/{TENANT_ID}/"
            )
            
            # 1. Validate Scope (Client permission)
            scopes = payload.get("scp", "").split()
            if "user_impersonation" not in scopes:
                raise HTTPException(status_code=403, detail="Invalid token scope.")

            # 2. Validate RBAC Roles (User permission)
            user_roles = payload.get("roles", []) # Entra App Roles array
            
            # Check if user has at least one of the required roles
            if not any(role in user_roles for role in self.required_roles):
                raise HTTPException(
                    status_code=status.HTTP_403_FORBIDDEN,
                    detail="User does not have the required application role."
                )
                
            return payload
            
        except jwt.ExpiredSignatureError:
            raise HTTPException(status_code=401, detail="Token expired")
        except jwt.InvalidTokenError as e:
            raise HTTPException(status_code=401, detail=f"Invalid token: {str(e)}")

app = FastAPI()

@app.get("/admin/reports")
def get_reports(current_user = Depends(EntraRoleChecker(["API.Admin"]))):
    return {"data": "Sensitive Admin Reports", "authenticated_as": current_user.get("name")}