Dynavera/apps/onboarding/viewsets.py

from django.db.models import Q
from django.db import transaction
from django.utils import timezone
from rest_framework.status import HTTP_400_BAD_REQUEST, HTTP_403_FORBIDDEN, HTTP_201_CREATED, HTTP_200_OK
from rest_framework.decorators import action
from rest_framework.permissions import IsAuthenticated
from rest_framework.response import Response
from rest_framework.viewsets import ModelViewSet, ReadOnlyModelViewSet

from apps.onboarding.models import AgentConfig, OnboardingSession, AgentInteractionLog, OnboardingFlow
from apps.onboarding.serializers import (
    AgentConfigSerializer, 
    OnboardingSessionSerializer, 
    AgentInteractionLogSerializer, 
    OnboardingFlowSerializer
)

class OnboardingFlowViewSet(ModelViewSet):
    queryset = OnboardingFlow.objects.all()
    serializer_class = OnboardingFlowSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        user = self.request.user
        queryset = OnboardingFlow.objects.filter(
            Q(role__organization__owner=user) | 
            Q(role__organization__members=user)
        ).distinct()

        role_uuid = self.request.query_params.get('role')
        if role_uuid:
            queryset = queryset.filter(role__uuid=role_uuid)

        return queryset.order_by('-created_at')

    def destroy(self, request, *args, **kwargs):
        flow = self.get_object()

        with transaction.atomic():
            OnboardingSession.objects.filter(role=flow.role).delete()
            self.perform_destroy(flow)

        return Response(status=204)

    @action(detail=True, methods=['post'], url_path='start-session')
    def start_session(self, request, uuid=None):
        flow = self.get_object()
        
        session, created = OnboardingSession.objects.get_or_create(
            user=request.user,
            role=flow.role,
            defaults={
                'status': 'active',
                'state': {
                    'progress': 0,
                    'current_step': 'intro',
                    'flow_uuid': str(flow.uuid),
                },
                'active_configs': {},
            }
        )

        if not created:
            state = session.state or {}
            state['flow_uuid'] = str(flow.uuid)
            session.state = state
            session.save(update_fields=['state', 'updated_at'])
        
        serializer = OnboardingSessionSerializer(session)
        return Response(serializer.data, status=HTTP_201_CREATED if created else HTTP_200_OK)

class AgentConfigViewSet(ModelViewSet):
    queryset = AgentConfig.objects.all()
    serializer_class = AgentConfigSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        return AgentConfig.objects.filter(organization__members=self.request.user).distinct()

    def perform_create(self, serializer):
        if not self.request.user.is_manager:
            return Response({'error': 'Forbidden'}, status=HTTP_403_FORBIDDEN)
        serializer.save()

class OnboardingSessionViewSet(ModelViewSet):
    queryset = OnboardingSession.objects.all()
    serializer_class = OnboardingSessionSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        user = self.request.user
        if user.is_manager:
            queryset = OnboardingSession.objects.filter(role__organization__members=user).distinct()
        else:
            queryset = OnboardingSession.objects.filter(user=user)

        role_uuid = self.request.query_params.get('role')
        if role_uuid:
            queryset = queryset.filter(role__uuid=role_uuid)

        return queryset.order_by('-created_at')

    @action(detail=True, methods=['post'], url_path='interact')
    def interact(self, request, uuid=None):
        session = self.get_object()
        user_message = request.data.get('message')
        page_uuid = request.data.get('page_uuid')
        responses = request.data.get('responses')

        if not user_message and not page_uuid:
            return Response({'error': 'Message or page_uuid is required'}, status=HTTP_400_BAD_REQUEST)

        if isinstance(responses, dict):
            state = session.state or {}
            stored_responses = state.get('responses', {})
            if not isinstance(stored_responses, dict):
                stored_responses = {}

            if page_uuid:
                stored_responses[str(page_uuid)] = responses
            else:
                stored_responses.update(responses)

            state['responses'] = stored_responses
            if page_uuid:
                state['last_page_uuid'] = str(page_uuid)
            session.state = state
            session.save(update_fields=['state', 'updated_at'])

        AgentInteractionLog.objects.create(
            session=session,
            sender_type='user',
            content=user_message or f'Submitted onboarding responses for page {page_uuid or "unknown"}',
            tool_call_metadata={'page_uuid': page_uuid, 'has_responses': isinstance(responses, dict)}
        )

                             
        
        return Response({
            'status': 'received',
            'session_state': session.state
        })

    @action(detail=True, methods=['get'], url_path='history')
    def history(self, request, uuid=None):
        session = self.get_object()
        logs = session.logs.all().order_by('created_at')
        serializer = AgentInteractionLogSerializer(logs, many=True)
        return Response(serializer.data)

    @action(detail=True, methods=['post'], url_path='complete')
    def complete(self, request, uuid=None):
        session = self.get_object()
        session.status = 'completed'
        session.completed_at = timezone.now()
        session.save()
        return Response({'message': 'Session marked as completed'})

class AgentInteractionLogViewSet(ReadOnlyModelViewSet):
    queryset = AgentInteractionLog.objects.all()
    serializer_class = AgentInteractionLogSerializer
    permission_classes = [IsAuthenticated]
    lookup_field = 'uuid'

    def get_queryset(self):
        return AgentInteractionLog.objects.filter(
            Q(session__user=self.request.user) | 
            Q(session__role__organization__owner=self.request.user)
        ).distinct()