AI AUtomtaion, Schudelign and publishign fromt and backe end refoactr

backend/igny8_core/modules/billing/views.py

@@ -840,3 +840,177 @@ class AIModelConfigViewSet(viewsets.ReadOnlyModelViewSet):
                 status_code=status.HTTP_404_NOT_FOUND
             )
 
+
+# ==============================================================================
+# Site AI Budget Allocation ViewSet
+# ==============================================================================
+
+from rest_framework import serializers as drf_serializers
+
+
+class SiteAIBudgetAllocationSerializer(drf_serializers.Serializer):
+    """Serializer for SiteAIBudgetAllocation model"""
+    id = drf_serializers.IntegerField(read_only=True)
+    ai_function = drf_serializers.CharField()
+    ai_function_display = drf_serializers.SerializerMethodField()
+    allocation_percentage = drf_serializers.IntegerField(min_value=0, max_value=100)
+    is_enabled = drf_serializers.BooleanField()
+    
+    def get_ai_function_display(self, obj):
+        display_map = {
+            'clustering': 'Keyword Clustering (Stage 1)',
+            'idea_generation': 'Ideas Generation (Stage 2)',
+            'content_generation': 'Content Generation (Stage 4)',
+            'image_prompt': 'Image Prompt Extraction (Stage 5)',
+            'image_generation': 'Image Generation (Stage 6)',
+        }
+        if hasattr(obj, 'ai_function'):
+            return display_map.get(obj.ai_function, obj.ai_function)
+        return display_map.get(obj.get('ai_function', ''), '')
+
+
+@extend_schema_view(
+    list=extend_schema(tags=['Billing'], summary='Get AI budget allocations for a site'),
+    create=extend_schema(tags=['Billing'], summary='Update AI budget allocations for a site'),
+)
+class SiteAIBudgetAllocationViewSet(viewsets.ViewSet):
+    """
+    ViewSet for managing Site AI Budget Allocations.
+    
+    GET /api/v1/billing/sites/{site_id}/ai-budget/
+    POST /api/v1/billing/sites/{site_id}/ai-budget/
+    
+    Allows configuring what percentage of the site's credit budget
+    can be used for each AI function during automation runs.
+    """
+    permission_classes = [IsAuthenticatedAndActive, HasTenantAccess]
+    authentication_classes = [JWTAuthentication]
+    throttle_scope = 'billing'
+    throttle_classes = [DebugScopedRateThrottle]
+    
+    def _get_site(self, site_id, request):
+        """Get site and verify user has access"""
+        from igny8_core.auth.models import Site
+        try:
+            site = Site.objects.get(id=int(site_id))
+            account = getattr(request, 'account', None)
+            if account and site.account != account:
+                return None
+            return site
+        except (Site.DoesNotExist, ValueError, TypeError):
+            return None
+    
+    def list(self, request, site_id=None):
+        """
+        Get AI budget allocations for a site.
+        Creates default allocations if they don't exist.
+        """
+        from igny8_core.business.billing.models import SiteAIBudgetAllocation
+        
+        site = self._get_site(site_id, request)
+        if not site:
+            return error_response(
+                message='Site not found or access denied',
+                errors=None,
+                status_code=status.HTTP_404_NOT_FOUND
+            )
+        
+        account = getattr(request, 'account', None) or site.account
+        
+        # Get or create default allocations
+        allocations = SiteAIBudgetAllocation.get_or_create_defaults_for_site(site, account)
+        
+        # Calculate total allocation
+        total_percentage = sum(a.allocation_percentage for a in allocations if a.is_enabled)
+        
+        serializer = SiteAIBudgetAllocationSerializer(allocations, many=True)
+        
+        return success_response(
+            data={
+                'allocations': serializer.data,
+                'total_percentage': total_percentage,
+                'is_valid': total_percentage <= 100,
+            },
+            message='AI budget allocations retrieved'
+        )
+    
+    def create(self, request, site_id=None):
+        """
+        Update AI budget allocations for a site.
+        
+        Body:
+        {
+            "allocations": [
+                {"ai_function": "clustering", "allocation_percentage": 15, "is_enabled": true},
+                {"ai_function": "idea_generation", "allocation_percentage": 10, "is_enabled": true},
+                {"ai_function": "content_generation", "allocation_percentage": 40, "is_enabled": true},
+                {"ai_function": "image_prompt", "allocation_percentage": 5, "is_enabled": true},
+                {"ai_function": "image_generation", "allocation_percentage": 30, "is_enabled": true}
+            ]
+        }
+        """
+        from igny8_core.business.billing.models import SiteAIBudgetAllocation
+        
+        site = self._get_site(site_id, request)
+        if not site:
+            return error_response(
+                message='Site not found or access denied',
+                errors=None,
+                status_code=status.HTTP_404_NOT_FOUND
+            )
+        
+        account = getattr(request, 'account', None) or site.account
+        allocations_data = request.data.get('allocations', [])
+        
+        if not allocations_data:
+            return error_response(
+                message='No allocations provided',
+                errors={'allocations': ['This field is required']},
+                status_code=status.HTTP_400_BAD_REQUEST
+            )
+        
+        # Validate total percentage
+        total_percentage = sum(
+            a.get('allocation_percentage', 0) 
+            for a in allocations_data 
+            if a.get('is_enabled', True)
+        )
+        
+        if total_percentage > 100:
+            return error_response(
+                message='Total allocation exceeds 100%',
+                errors={'total_percentage': [f'Total is {total_percentage}%, must be <= 100%']},
+                status_code=status.HTTP_400_BAD_REQUEST
+            )
+        
+        # Update or create allocations
+        valid_functions = ['clustering', 'idea_generation', 'content_generation', 'image_prompt', 'image_generation']
+        updated = []
+        
+        for alloc_data in allocations_data:
+            ai_function = alloc_data.get('ai_function')
+            if ai_function not in valid_functions:
+                continue
+            
+            allocation, _ = SiteAIBudgetAllocation.objects.update_or_create(
+                account=account,
+                site=site,
+                ai_function=ai_function,
+                defaults={
+                    'allocation_percentage': alloc_data.get('allocation_percentage', 20),
+                    'is_enabled': alloc_data.get('is_enabled', True),
+                }
+            )
+            updated.append(allocation)
+        
+        serializer = SiteAIBudgetAllocationSerializer(updated, many=True)
+        
+        return success_response(
+            data={
+                'allocations': serializer.data,
+                'total_percentage': total_percentage,
+            },
+            message='AI budget allocations updated successfully'
+        )
+
+