from rest_framework_simplejwt.views import TokenObtainPairView, TokenRefreshView
from rest_framework_simplejwt.serializers import TokenObtainPairSerializer
from .serializers import UserSerializer, CookieTokenRefreshSerializer
from rest_framework.response import Response
from rest_framework import status
from django.conf import settings
from django.utils import timezone
from rest_framework.views import APIView

class CustomTokenObtainPairSerializer(TokenObtainPairSerializer):
    @classmethod
    def get_token(cls, user):
        token = super().get_token(user)
        token["username"] = user.username
        token["email"] = user.email
        return token

class CustomTokenObtainPairView(TokenObtainPairView):
    serializer_class = CustomTokenObtainPairSerializer

    def post(self, request, *args, **kwargs):
        serializer = self.get_serializer(data=request.data)
        try:
            serializer.is_valid(raise_exception=True)
        except Exception:
            return super().post(request, *args, **kwargs)

        access_token = serializer.validated_data.get('access')
        refresh_token = serializer.validated_data.get('refresh')

        user_serializer = UserSerializer(serializer.user)

        response = Response({'user': user_serializer.data}, status=status.HTTP_200_OK)

        access_token_expiration = timezone.now() + settings.SIMPLE_JWT['ACCESS_TOKEN_LIFETIME']
        refresh_token_expiration = timezone.now() + settings.SIMPLE_JWT['REFRESH_TOKEN_LIFETIME']

        response.set_cookie(
            'access_token',
            access_token,
            expires=access_token_expiration,
            httponly=True,
            samesite='Lax',
            secure=not settings.DEBUG
        )
        response.set_cookie(
            'refresh_token',
            refresh_token,
            expires=refresh_token_expiration,
            httponly=True,
            samesite='Lax',
            secure=not settings.DEBUG
        )
        return response

class CustomTokenRefreshView(TokenRefreshView):
    serializer_class = CookieTokenRefreshSerializer

    def post(self, request, *args, **kwargs):
        response = super().post(request, *args, **kwargs)
        if response.status_code == 200:
            access_token = response.data.pop('access')
            response.set_cookie(
                'access_token',
                access_token,
                httponly=True,
                samesite='Lax',
                secure=not settings.DEBUG
            )
            if 'refresh' in response.data:
                refresh_token = response.data.pop('refresh')
                response.set_cookie(
                    'refresh_token',
                    refresh_token,
                    httponly=True,
                    samesite='Lax',
                    secure=not settings.DEBUG
                )
        return response

class LogoutView(APIView):
    def post(self, request, *args, **kwargs):
        response = Response(status=status.HTTP_204_NO_CONTENT)
        response.delete_cookie('access_token')
        response.delete_cookie('refresh_token')
        return response