from typing import Union
from functools import wraps

from fastapi import HTTPException, Request

from .models import User


def get_current_user(request: Request) -> User:
    auth_data = request.session.get("user")
    if auth_data is None:
        raise HTTPException(401, "Not authorized")
    return User(**auth_data)


def get_current_user_or_none(request: Request) -> User | None:
    try:
        return get_current_user(request)
    except HTTPException:
        return None


def hasrole(
    required_roles: Union[str, list[str]] = [],
    roles_key: str = "roles",
    realm: str | None = "realm_access",  # Keycloak standard for realm defined roles
):
    required_roles_set: set[str]
    if isinstance(required_roles, str):
        required_roles_set = set([required_roles])
    else:
        required_roles_set = set(required_roles)

    def decorator(func):
        @wraps(func)
        async def wrapper(request=None, *args, **kwargs):
            if request is None:
                raise HTTPException(
                    500,
                    "Functions decorated with hasrole must have a request:Request argument",
                )
            if "user" not in request.session:
                raise HTTPException(401, "Not authorized")
            user = request.session["user"]
            try:
                if realm in user:
                    roles = user[realm][roles_key]
                else:
                    roles = user[roles_key]
            except KeyError:
                raise HTTPException(401, "Not authorized")
            if not any(required_roles_set.intersection(roles)):
                raise HTTPException(401, "Not authorized")
            return await func(request, *args, **kwargs)

        return wrapper

    return decorator