# @title utils/jwt_handler.py
import json
import base64
from typing import Optional, Dict, Any

class JWTHandler:
    """
    Klasa do obsługi tokenów JWT i ekstrakcji GroupID.
    """
    @staticmethod
    def decode_jwt(jwt_token: str) -> Optional[Dict[str, Any]]:
        """
        Dekoduje token JWT i zwraca payload.
        
        Args:
            jwt_token: Token JWT do zdekodowania
            
        Returns:
            Dict z danymi z payloadu lub None w przypadku błędu
        """
        try:
            # Token JWT składa się z 3 części oddzielonych kropkami
            header, payload, signature = jwt_token.split('.')
            
            # Dekoduj payload
            # Uzupełnij padding jeśli potrzebny
            padding = 4 - (len(payload) % 4)
            if padding != 4:
                payload = payload + ('=' * padding)
                
            # Zamień znaki specjalne
            payload = payload.replace('-', '+').replace('_', '/')
            
            # Dekoduj Base64
            decoded_bytes = base64.b64decode(payload)
            decoded_str = decoded_bytes.decode('utf-8')
            
            # Parsuj JSON
            payload_data = json.loads(decoded_str)
            
            return payload_data
            
        except Exception as e:
            print(f"Błąd dekodowania JWT: {str(e)}")
            return None
    
    @staticmethod
    def extract_group_id(jwt_token: str) -> Optional[str]:
        """
        Ekstrahuje GroupID z tokenu JWT.
        
        Args:
            jwt_token: Token JWT
            
        Returns:
            GroupID lub None w przypadku błędu
        """
        try:
            payload = JWTHandler.decode_jwt(jwt_token)
            if payload and 'GroupID' in payload:
                return payload['GroupID']
            # Sprawdź alternatywne pole
            elif payload and 'GroupId' in payload:
                return payload['GroupId']
            return None
        except Exception as e:
            print(f"Błąd ekstrakcji GroupID: {str(e)}")
            return None
    
    @staticmethod
    def validate_jwt(jwt_token: str) -> bool:
        """
        Sprawdza czy token JWT jest poprawny i zawiera GroupID.
        
        Args:
            jwt_token: Token JWT do walidacji
            
        Returns:
            True jeśli token jest poprawny i zawiera GroupID
        """
        try:
            payload = JWTHandler.decode_jwt(jwt_token)
            return payload is not None and 'GroupID' in payload
        except:
            return False