bookshop/simulations/views.py

from rest_framework import viewsets, status
from rest_framework.decorators import action
from rest_framework.response import Response
from .models import SimulationMethod, SimulationRun
from .serializers import (
    SimulationMethodSerializer,
    SimulationRunSerializer,
    SimulationRunCreateSerializer,
)
from .tasks import run_simulation


class SimulationMethodViewSet(viewsets.ModelViewSet):
    queryset = SimulationMethod.objects.filter(is_active=True)
    serializer_class = SimulationMethodSerializer
    lookup_field = 'slug'


class SimulationRunViewSet(viewsets.ModelViewSet):
    queryset = SimulationRun.objects.all()
    
    def get_serializer_class(self):
        if self.action == 'create':
            return SimulationRunCreateSerializer
        if self.action == 'retrieve':
            return SimulationRunDetailSerializer
        return SimulationRunSerializer
    
    def create(self, request, *args, **kwargs):
        serializer = self.get_serializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        
        run = serializer.save()
        
        run_simulation.delay(run.id)
        
        output_serializer = SimulationRunSerializer(run)
        return Response(
            output_serializer.data,
            status=status.HTTP_202_ACCEPTED
        )
    
    @action(detail=True, methods=['post'])
    def cancel(self, request, pk=None):
        run = self.get_object()
        if run.status in ['PENDING', 'RUNNING']:
            run.status = 'CANCELLED'
            run.save(update_fields=['status'])
            return Response({'status': 'cancelled'})
        return Response(
            {'error': 'Cannot cancel run in current state'},
            status=status.HTTP_400_BAD_REQUEST
        )
    
    @action(detail=True, methods=['get'])
    def status(self, request, pk=None):
        run = self.get_object()
        return Response({
            'id': run.id,
            'status': run.status,
            'progress': run.progress,
            'error_message': run.error_message,
            'created_at': run.created_at,
            'started_at': run.started_at,
            'completed_at': run.completed_at,
        })
    
    @action(detail=False, methods=['get'])
    def pending(self, request):
        runs = SimulationRun.objects.filter(status='PENDING')
        serializer = self.get_serializer(runs, many=True)
        return Response(serializer.data)
    
    @action(detail=False, methods=['get'])
    def running(self, request):
        runs = SimulationRun.objects.filter(status='RUNNING')
        serializer = self.get_serializer(runs, many=True)
        return Response(serializer.data)
    
    @action(detail=False, methods=['get'])
    def recent(self, request):
        limit = int(request.query_params.get('limit', 10))
        runs = SimulationRun.objects.all()[:limit]
        serializer = self.get_serializer(runs, many=True)
        return Response(serializer.data)