#!/usr/bin/env python3
"""
AI Providers Configuration for Lifestyle Journey Application

This module defines configurations for different AI providers (Google Gemini, Anthropic Claude)
and maps specific agents to their preferred providers and models.
"""

import os
from typing import Dict, Any, Optional
from enum import Enum

class AIProvider(Enum):
    """Supported AI providers"""
    GEMINI = "gemini"
    ANTHROPIC = "anthropic"

class AIModel(Enum):
    """Supported AI models"""
    # Gemini models
    GEMINI_2_5_FLASH = "gemini-2.5-flash"
    GEMINI_2_0_FLASH = "gemini-2.0-flash"
    GEMINI_2_5_PRO = "gemini-2.5-pro"
    GEMINI_1_5_PRO = "gemini-1.5-pro"
    
    # Anthropic models
    CLAUDE_SONNET_4 = "claude-sonnet-4-20250514"
    CLAUDE_SONNET_3_7 = "claude-3-7-sonnet-20250219"
    CLAUDE_SONNET_3_5 = "claude-3-5-sonnet-20241022"
    CLAUDE_HAIKU_3_5 = "claude-3-5-haiku-20241022"

# Provider-specific configurations
PROVIDER_CONFIGS = {
    AIProvider.GEMINI: {
        "api_key_env": "GEMINI_API_KEY",
        "default_model": AIModel.GEMINI_2_0_FLASH,
        "default_temperature": 0.3,
        "max_tokens": None,  # Gemini handles this automatically
        "available_models": [
            AIModel.GEMINI_2_5_FLASH,
            AIModel.GEMINI_2_0_FLASH,
            AIModel.GEMINI_2_5_PRO,
            AIModel.GEMINI_1_5_PRO
        ]
    },
    AIProvider.ANTHROPIC: {
        "api_key_env": "ANTHROPIC_API_KEY",
        "default_model": AIModel.CLAUDE_SONNET_4,
        "default_temperature": 0.3,
        "max_tokens": 20000,
        "available_models": [
            AIModel.CLAUDE_SONNET_4,
            AIModel.CLAUDE_SONNET_3_7,
            AIModel.CLAUDE_SONNET_3_5,
            AIModel.CLAUDE_HAIKU_3_5
        ]
    }
}

# Agent-specific provider and model assignments
AGENT_CONFIGURATIONS = {
    # Main Lifestyle Assistant uses Anthropic Claude
    "MainLifestyleAssistant": {
        "provider": AIProvider.ANTHROPIC,
        "model": AIModel.CLAUDE_SONNET_4,
        "temperature": 0.2,
        "reasoning": "Complex lifestyle coaching requires advanced reasoning capabilities"
    },
    
    # All other agents use Google Gemini
    "EntryClassifier": {
        "provider": AIProvider.GEMINI,
        "model": AIModel.GEMINI_2_0_FLASH,
        "temperature": 0.1,
        "reasoning": "Fast classification task, optimized for speed"
    },
    
    "TriageExitClassifier": {
        "provider": AIProvider.GEMINI,
        "model": AIModel.GEMINI_2_0_FLASH,
        "temperature": 0.2,
        "reasoning": "Medical triage decisions require consistency"
    },
    
    "MedicalAssistant": {
        "provider": AIProvider.ANTHROPIC,
        "model": AIModel.CLAUDE_SONNET_4,
        "temperature": 0.2,
        "reasoning": "Medical guidance requires reliable, consistent responses"
    },
    
    "SoftMedicalTriage": {
        "provider": AIProvider.GEMINI,
        "model": AIModel.GEMINI_2_0_FLASH,
        "temperature": 0.3,
        "reasoning": "Gentle triage can use faster model"
    },
    
    "LifestyleProfileUpdater": {
        "provider": AIProvider.GEMINI,
        "model": AIModel.GEMINI_2_5_FLASH,
        "temperature": 0.2,
        "reasoning": "Profile analysis requires detailed processing"
    }
}

def get_agent_config(agent_name: str) -> Dict[str, Any]:
    """
    Get configuration for a specific agent
    
    Args:
        agent_name: Name of the agent (e.g., "MainLifestyleAssistant")
        
    Returns:
        Dictionary with provider, model, and other configuration details
    """
    if agent_name not in AGENT_CONFIGURATIONS:
        # Default to Gemini for unknown agents
        return {
            "provider": AIProvider.GEMINI,
            "model": AIModel.GEMINI_2_5_FLASH,
            "temperature": 0.3,
            "reasoning": "Default configuration for unknown agent"
        }
    
    return AGENT_CONFIGURATIONS[agent_name].copy()

def get_provider_config(provider: AIProvider) -> Dict[str, Any]:
    """
    Get configuration for a specific provider
    
    Args:
        provider: AI provider enum
        
    Returns:
        Dictionary with provider-specific configuration
    """
    return PROVIDER_CONFIGS[provider].copy()

def is_provider_available(provider: AIProvider) -> bool:
    """
    Check if a provider is available (has API key configured)
    
    Args:
        provider: AI provider to check
        
    Returns:
        True if provider is available, False otherwise
    """
    config = get_provider_config(provider)
    api_key = os.getenv(config["api_key_env"])
    return api_key is not None and api_key.strip() != ""

def get_available_providers() -> list[AIProvider]:
    """
    Get list of available providers (those with API keys configured)
    
    Returns:
        List of available AI providers
    """
    available = []
    for provider in AIProvider:
        if is_provider_available(provider):
            available.append(provider)
    return available

def validate_configuration() -> Dict[str, Any]:
    """
    Validate the current AI provider configuration
    
    Returns:
        Dictionary with validation results
    """
    results = {
        "valid": True,
        "errors": [],
        "warnings": [],
        "available_providers": [],
        "agent_status": {}
    }
    
    # Check available providers
    available_providers = get_available_providers()
    results["available_providers"] = [p.value for p in available_providers]
    
    if not available_providers:
        results["valid"] = False
        results["errors"].append("No AI providers available - check API keys")
        return results
    
    # Check each agent configuration
    for agent_name, config in AGENT_CONFIGURATIONS.items():
        provider = config["provider"]
        model = config["model"]
        
        agent_status = {
            "provider": provider.value,
            "model": model.value,
            "available": provider in available_providers,
            "fallback_needed": False
        }
        
        if provider not in available_providers:
            agent_status["fallback_needed"] = True
            results["warnings"].append(
                f"Agent {agent_name} configured for {provider.value} but provider not available"
            )
            
            # Suggest fallback
            if AIProvider.GEMINI in available_providers:
                agent_status["fallback_provider"] = AIProvider.GEMINI.value
                agent_status["fallback_model"] = AIModel.GEMINI_2_5_FLASH.value
            elif available_providers:
                fallback = available_providers[0]
                agent_status["fallback_provider"] = fallback.value
                fallback_config = get_provider_config(fallback)
                agent_status["fallback_model"] = fallback_config["default_model"].value
        
        results["agent_status"][agent_name] = agent_status
    
    return results

# Environment variable validation
def check_environment_setup() -> Dict[str, str]:
    """
    Check which AI provider API keys are configured
    
    Returns:
        Dictionary mapping provider names to their status
    """
    status = {}
    
    for provider in AIProvider:
        config = get_provider_config(provider)
        api_key_env = config["api_key_env"]
        api_key = os.getenv(api_key_env)
        
        if api_key and api_key.strip():
            status[provider.value] = "✅ Configured"
        else:
            status[provider.value] = f"❌ Missing {api_key_env}"
    
    return status

if __name__ == "__main__":
    print("🤖 AI Providers Configuration")
    print("=" * 50)
    
    # Check environment setup
    print("\n📋 Environment Setup:")
    env_status = check_environment_setup()
    for provider, status in env_status.items():
        print(f"   {provider}: {status}")
    
    # Validate configuration
    print("\n🔍 Configuration Validation:")
    validation = validate_configuration()
    
    if validation["valid"]:
        print("   ✅ Configuration is valid")
    else:
        print("   ❌ Configuration has errors:")
        for error in validation["errors"]:
            print(f"      - {error}")
    
    if validation["warnings"]:
        print("   ⚠️ Warnings:")
        for warning in validation["warnings"]:
            print(f"      - {warning}")
    
    print(f"\n📊 Available Providers: {', '.join(validation['available_providers'])}")
    
    print("\n🎯 Agent Assignments:")
    for agent, status in validation["agent_status"].items():
        provider_info = f"{status['provider']} ({status['model']})"
        availability = "✅" if status["available"] else "❌"
        print(f"   {agent}: {provider_info} {availability}")
        
        if status.get("fallback_needed"):
            fallback_info = f"{status.get('fallback_provider')} ({status.get('fallback_model')})"
            print(f"      → Fallback: {fallback_info}")