import jwt
from fastapi_users.authentication import JWTStrategy
from fastapi_users.jwt import generate_jwt, decode_jwt
from fastapi_users import exceptions
from app.config import SECRET


class CustomJWTStrategy(JWTStrategy):
    async def write_token(self, user) -> str:
        # Coletar todas as permissões do usuário a partir de seus papéis
        todas_as_permissoes = set()  # Usamos um set para evitar duplicatas

        if hasattr(user, 'papeis'):
            for papel in user.papeis:
                if hasattr(papel, 'permissoes'):
                    for permissao in papel.permissoes:
                        todas_as_permissoes.add(permissao.id)  # Usar o ID da permissão

        # Transformar o set em uma lista para o payload do token
        lista_de_permissoes = list(todas_as_permissoes)
        print("user id")
        print(user.id)

        # Aqui, adicionamos as claims personalizadas ao payload
        data = {
            "sub": str(user.id),
            "permissions": lista_de_permissoes,  # Acessa diretamente a lista de IDs de permissões coletadas
            "aud": self.token_audience,  # Audiência, conforme o padrão
        }
        token = generate_jwt(
            data, self.encode_key, self.lifetime_seconds, algorithm=self.algorithm
        )

        return token

    async def read_token(self, token: str, user_manager):
        # Decodifica o token JWT usando a função padrão decode_jwt do fastapi_users
        try:
            payload = decode_jwt(token, SECRET, audience=self.token_audience)
            return payload
        except Exception as e:
            raise ValueError(f"Token inválido: {str(e)}")