from django.shortcuts import get_object_or_404
from django.utils import timezone
from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework.parsers import MultiPartParser, FormParser
from .models import Client, Task, TaskResult
from .serializers import (
    ClientSerializer,
    TaskSerializer,
    TaskResultSerializer,
    TaskClaimSerializer,
    TaskStartSerializer,
    TaskCompleteSerializer
)

class ClientViewSet(viewsets.ModelViewSet):
    queryset = Client.objects.all()
    serializer_class = ClientSerializer
    permission_classes = [IsAuthenticated]

class TaskViewSet(viewsets.ModelViewSet):
    queryset = Task.objects.all()
    serializer_class = TaskSerializer
    permission_classes = [IsAuthenticated]

    @action(detail=False, methods=['post'], serializer_class=TaskClaimSerializer)
    def claim(self, request):
        """Client claims an available task"""
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        client_name = serializer.validated_data['client_name']

        # Try to find a pending task, either assigned to this client or unassigned
        with self.queryset.model.objects._base_manager._db.transaction.atomic():
            # First try to find tasks assigned to this client
            task = Task.objects.filter(
                status='pending',
                client_name=client_name
            ).select_for_update().first()

            # If no task assigned to this client, try to find unassigned tasks
            if not task:
                task = Task.objects.filter(
                    status='pending',
                    client_name__isnull=True
                ).select_for_update().first()

            if not task:
                return Response({'detail': 'No available tasks'}, status=status.HTTP_404_NOT_FOUND)

            # Assign the task to the client
            task.status = 'assigned'
            task.assigned_to = client_name
            task.save()

        return Response(TaskSerializer(task).data, status=status.HTTP_200_OK)

    @action(detail=True, methods=['post'], serializer_class=TaskStartSerializer)
    def start(self, request, pk=None):
        """Client starts a task"""
        task = self.get_object()
        if task.status != 'assigned':
            return Response({'detail': 'Task is not assigned'}, status=status.HTTP_400_BAD_REQUEST)

        task.status = 'running'
        task.started_at = timezone.now()
        task.save()

        return Response(TaskSerializer(task).data, status=status.HTTP_200_OK)

    @action(detail=True, methods=['post'], serializer_class=TaskCompleteSerializer)
    def complete(self, request, pk=None):
        """Client completes a task"""
        task = self.get_object()
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)

        if task.status != 'running':
            return Response({'detail': 'Task is not running'}, status=status.HTTP_400_BAD_REQUEST)

        task.status = serializer.validated_data['status']
        task.completed_at = timezone.now()
        task.save()

        return Response(TaskSerializer(task).data, status=status.HTTP_200_OK)

class TaskResultViewSet(viewsets.ModelViewSet):
    queryset = TaskResult.objects.all()
    serializer_class = TaskResultSerializer
    permission_classes = [IsAuthenticated]
    parser_classes = [MultiPartParser, FormParser]

    @action(detail=True, methods=['get'])
    def download(self, request, pk=None):
        """Download task result file"""
        task_result = self.get_object()
        if not task_result.result_file:
            return Response({'detail': 'No file available'}, status=status.HTTP_404_NOT_FOUND)

        # Use Django's built-in FileResponse for downloading
        from django.http import FileResponse
        return FileResponse(task_result.result_file.open('rb'), as_attachment=True)