from rest_framework_simplejwt.authentication import JWTAuthentication
from rest_framework_simplejwt.exceptions import InvalidToken, AuthenticationFailed
from django.conf import settings
from rest_framework.authentication import CSRFCheck
from rest_framework import exceptions

def enforce_csrf(request):
    """
    Enforce CSRF validation for cookie-based authentication.
    """
    # Dummy get_response for CSRFCheck initialization
    def dummy_get_response(request):
        return None
        
    check = CSRFCheck(dummy_get_response)
    check.process_request(request)
    reason = check.process_view(request, None, (), {})
    if reason:
        raise exceptions.PermissionDenied('CSRF Failed: %s' % reason)

class CookieJWTAuthentication(JWTAuthentication):
    """
    Custom authentication class that checks for a JWT token in the cookies
    if it's not provided in the Authorization header.
    """
    def authenticate(self, request):
        # 1. Try to get the token from the header (standard behavior)
        header = self.get_header(request)
        
        if header is None:
            # 2. If no header, try to get the token from the cookie
            cookie_name = settings.SIMPLE_JWT.get('AUTH_COOKIE', 'access_token')
            raw_token = request.COOKIES.get(cookie_name)
        else:
            raw_token = self.get_raw_token(header)
            
        if raw_token is None:
            return None

        try:
            validated_token = self.get_validated_token(raw_token)
            
            # Enforce CSRF if token was provided via cookie
            if header is None:
                enforce_csrf(request)
                
            return self.get_user(validated_token), validated_token
        except (InvalidToken, AuthenticationFailed):
            return None