from rest_framework_simplejwt.views import TokenObtainPairView, TokenRefreshView
from rest_framework.response import Response
from rest_framework import status
from django.conf import settings

class CookieTokenObtainPairView(TokenObtainPairView):
    """
    Custom view to set JWT tokens as HttpOnly cookies.
    """
    def post(self, request, *args, **kwargs):
        response = super().post(request, *args, **kwargs)
        if response.status_code == 200:
            access = response.data.get('access')
            refresh = response.data.get('refresh')
            # Remove tokens from body (optional, for extra security)
            response.data.pop('access', None)
            response.data.pop('refresh', None)
            # Set cookies
            cookie_settings = {
                'httponly': True,
                'secure': True,  # Set to True if using HTTPS
                'samesite': 'Lax',
                'path': '/'
            }
            response.set_cookie('access_token', access, max_age=60*5, **cookie_settings)  # 5 min
            response.set_cookie('refresh_token', refresh, max_age=60*60*24*7, **cookie_settings)  # 7 days
        return response

class CookieTokenRefreshView(TokenRefreshView):
    """
    Custom view to refresh JWT tokens and set as HttpOnly cookies.
    """
    def post(self, request, *args, **kwargs):
        response = super().post(request, *args, **kwargs)
        if response.status_code == 200:
            access = response.data.get('access')
            response.data.pop('access', None)
            cookie_settings = {
                'httponly': True,
                'secure': True,  # Set to True if using HTTPS
                'samesite': 'Lax',
                'path': '/'
            }
            response.set_cookie('access_token', access, max_age=60*5, **cookie_settings)  # 5 min
        return response