ALwrity/backend/services/subscription/preflight_validator.py

"""
Pre-flight Validation Utility for Multi-Operation Workflows

Provides transparent validation for operations that involve multiple API calls.
Services can use this to validate entire workflows before making any external API calls.
"""

from typing import Dict, Any, List, Optional, Tuple
from fastapi import HTTPException
from loguru import logger

from services.subscription.pricing_service import PricingService
from models.subscription_models import APIProvider


def validate_research_operations(
    pricing_service: PricingService,
    user_id: str,
    gpt_provider: str = "google"
) -> None:
    """
    Validate all operations for a research workflow before making ANY API calls.
    
    This prevents wasteful external API calls (like Google Grounding) if subsequent
    LLM calls would fail due to token or call limits.
    
    Args:
        pricing_service: PricingService instance
        user_id: User ID for subscription checking
        gpt_provider: GPT provider from env var (defaults to "google")
        
    Returns:
        (can_proceed, error_message, error_details)
        If can_proceed is False, raises HTTPException with 429 status
    """
    try:
        # Determine actual provider for LLM calls based on GPT_PROVIDER env var
        gpt_provider_lower = gpt_provider.lower()
        if gpt_provider_lower == "huggingface":
            llm_provider_enum = APIProvider.MISTRAL  # Maps to HuggingFace
            llm_provider_name = "huggingface"
        else:
            llm_provider_enum = APIProvider.GEMINI
            llm_provider_name = "gemini"
        
        # Estimate tokens for each operation in research workflow
        # Google Grounding call: ~2000 tokens (input + output)
        # Keyword analyzer: ~1000 tokens (input: 3000 chars research, output: structured JSON)
        # Competitor analyzer: ~1000 tokens (input: 3000 chars research, output: structured JSON)
        # Content angle generator: ~1000 tokens (input: 3000 chars research, output: list of angles)
        
        operations_to_validate = [
            {
                'provider': APIProvider.GEMINI,  # Google Grounding uses Gemini
                'tokens_requested': 2000,
                'actual_provider_name': 'gemini',
                'operation_type': 'google_grounding'
            },
            {
                'provider': llm_provider_enum,
                'tokens_requested': 1000,
                'actual_provider_name': llm_provider_name,
                'operation_type': 'keyword_analysis'
            },
            {
                'provider': llm_provider_enum,
                'tokens_requested': 1000,
                'actual_provider_name': llm_provider_name,
                'operation_type': 'competitor_analysis'
            },
            {
                'provider': llm_provider_enum,
                'tokens_requested': 1000,
                'actual_provider_name': llm_provider_name,
                'operation_type': 'content_angle_generation'
            }
        ]
        
        logger.info(f"[Pre-flight Validator] 🚀 Starting Research Workflow Validation")
        logger.info(f"   ├─ User: {user_id}")
        logger.info(f"   ├─ LLM Provider: {llm_provider_name} (GPT_PROVIDER={gpt_provider})")
        logger.info(f"   └─ Operations to validate: {len(operations_to_validate)}")
        
        can_proceed, message, error_details = pricing_service.check_comprehensive_limits(
            user_id=user_id,
            operations=operations_to_validate
        )
        
        if not can_proceed:
            usage_info = error_details.get('usage_info', {}) if error_details else {}
            provider = usage_info.get('provider', llm_provider_name) if usage_info else llm_provider_name
            operation_type = usage_info.get('operation_type', 'unknown')
            
            logger.error(f"[Pre-flight Validator] ❌ RESEARCH WORKFLOW BLOCKED")
            logger.error(f"   ├─ User: {user_id}")
            logger.error(f"   ├─ Blocked at: {operation_type}")
            logger.error(f"   ├─ Provider: {provider}")
            logger.error(f"   └─ Reason: {message}")
            
            # Raise HTTPException immediately - frontend gets immediate response, no API calls made
            raise HTTPException(
                status_code=429,
                detail={
                    'error': message,
                    'message': message,
                    'provider': provider,
                    'usage_info': usage_info if usage_info else error_details
                }
            )
        
        logger.info(f"[Pre-flight Validator] ✅ RESEARCH WORKFLOW APPROVED")
        logger.info(f"   ├─ User: {user_id}")
        logger.info(f"   └─ All {len(operations_to_validate)} operations validated - proceeding with API calls")
        # Validation passed - no return needed (function raises HTTPException if validation fails)
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"[Pre-flight Validator] Error validating research operations: {e}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail={
                'error': f"Failed to validate operations: {str(e)}",
                'message': f"Failed to validate operations: {str(e)}"
            }
        )


def validate_image_generation_operations(
    pricing_service: PricingService,
    user_id: str
) -> None:
    """
    Validate image generation operation before making API calls.
    
    Args:
        pricing_service: PricingService instance
        user_id: User ID for subscription checking
        
    Returns:
        (can_proceed, error_message, error_details)
        If can_proceed is False, raises HTTPException with 429 status
    """
    try:
        operations_to_validate = [
            {
                'provider': APIProvider.STABILITY,
                'tokens_requested': 0,
                'actual_provider_name': 'stability',
                'operation_type': 'image_generation'
            }
        ]
        
        can_proceed, message, error_details = pricing_service.check_comprehensive_limits(
            user_id=user_id,
            operations=operations_to_validate
        )
        
        if not can_proceed:
            logger.error(f"[Pre-flight Validator] Image generation blocked for user {user_id}: {message}")
            
            usage_info = error_details.get('usage_info', {}) if error_details else {}
            provider = usage_info.get('provider', 'stability') if usage_info else 'stability'
            
            raise HTTPException(
                status_code=429,
                detail={
                    'error': message,
                    'message': message,
                    'provider': provider,
                    'usage_info': usage_info if usage_info else error_details
                }
            )
        
        logger.info(f"[Pre-flight Validator] ✅ Image generation validated for user {user_id}")
        # Validation passed - no return needed (function raises HTTPException if validation fails)
        
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"[Pre-flight Validator] Error validating image generation: {e}", exc_info=True)
        raise HTTPException(
            status_code=500,
            detail={
                'error': f"Failed to validate image generation: {str(e)}",
                'message': f"Failed to validate image generation: {str(e)}"
            }
        )