import logging
from typing import Optional, Dict, Any
import services.ai_service_legacy as legacy_ai
import services.ai_service_enterprise as enterprise_ai

logger = logging.getLogger("chatbot.ai_router_manager")

async def generate_response(
    query: str, 
    context: str, 
    custom_prompt: Optional[str] = None,
    use_enterprise: bool = False
) -> Dict[str, Any]:
    """
    Safely routes AI requests.
    - If `use_enterprise` is True (e.g. called from new SaaS endpoints), 
      it hits the Local->Cloud fallback chain.
    - Otherwise, it routes to the existing legacy Gemini->OpenAI chain, 
      ensuring zero disruption to current frontend widgets.
    """
    if use_enterprise:
        logger.info("Routing request to Enterprise AI Layer")
        return await enterprise_ai.generate_chatbot_response(query, context, custom_prompt)
    else:
        logger.info("Routing request to Legacy AI Layer")
        return await legacy_ai.generate_chatbot_response(query, context, custom_prompt)

async def stream_response(
    query: str, 
    context: str, 
    custom_prompt: Optional[str] = None,
    use_enterprise: bool = False
):
    """
    Routes streaming responses safely.
    Note: Enterprise streaming is not fully implemented in the enterprise module yet, 
    so it falls back to legacy streaming if enterprise streaming is missing.
    """
    if use_enterprise and hasattr(enterprise_ai, 'stream_chatbot_response'):
        logger.info("Routing stream to Enterprise AI Layer")
        async for chunk in enterprise_ai.stream_chatbot_response(query, context, custom_prompt):
            yield chunk
    else:
        logger.info("Routing stream to Legacy AI Layer")
        async for chunk in legacy_ai.stream_chatbot_response(query, context, custom_prompt):
            yield chunk