from typing import Union, Dict, Any

from django.apps import apps
from supertokens_python.recipe.session.interfaces import RecipeInterface
from supertokens_python.recipe.session.asyncio import get_session_information


def override_functions(original_implementation: RecipeInterface):
    original_implementation_update_access_token_payload = original_implementation.update_access_token_payload

    async def update_access_token_payload(session_handle: str,
                                          new_access_token_payload: Dict[str, Any],
                                          user_context: Dict[str, Any]) -> None:
        profile_role_model = apps.get_model('profile', 'ProfileRole')

        user_info = await get_session_information(session_handle)

        if new_access_token_payload is None:
            new_access_token_payload = {}

        new_access_token_payload["roles"] = {}
        profile_roles = profile_role_model.objects.filter(profile__supertokens_user_id=user_info.user_id)
        for profile_role in profile_roles:
            if profile_role.at not in new_access_token_payload["roles"]:
                new_access_token_payload["roles"][profile_role.at] = [profile_role.role]
            else:
                new_access_token_payload["roles"][profile_role.at].append(profile_role.role)

        return await original_implementation_update_access_token_payload(session_handle,
                                                                         new_access_token_payload, user_context)

    original_implementation.update_access_token_payload = update_access_token_payload
    return original_implementation