src/model_config.py

"""
Model configuration and management for personalization features.
Supports Groq (primary content generation), HuggingFace (visualization), and Gemini (TTS).
"""

import os
import logging
from typing import Optional
from openai import OpenAI
from anthropic import Anthropic
from dotenv import load_dotenv

load_dotenv()
logger = logging.getLogger(__name__)

# ============================================================================
# ENVIRONMENT VARIABLES
# ============================================================================

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
ANTHROPIC_API_KEY = os.getenv("ANTHROPIC_API_KEY")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# Groq models for content generation (fast, free tier available)
GROQ_CONTENT_MODELS = [
    "llama-3.3-70b-versatile",  # Current primary - excellent quality
    "llama-3.1-70b-versatile",  # Stable alternative
    "mixtral-8x7b-32768",       # Deprecated - kept for reference
]

USING_OPENAI = bool(OPENAI_API_KEY)
USING_ANTHROPIC = bool(ANTHROPIC_API_KEY)
USING_GROQ = bool(GROQ_API_KEY)
OPENAI_CLIENT = None
ANTHROPIC_CLIENT = None
GROQ_CLIENT = None

# Per-task primary model names
OPENAI_CONTENT_MODEL = "gpt-4o-mini"
ANTHROPIC_VISUALIZATION_MODEL = "claude-sonnet-4-20250514"
GROQ_CONTENT_MODEL = GROQ_CONTENT_MODELS[0]

# Generic placeholders
LLM_CLIENT = None
PRIMARY_MODEL = None
LLM_SERVICE = None

# LLM CLIENT INITIALIZATION
if USING_OPENAI:
    logger.info(f"OPENAI_API_KEY present: {bool(OPENAI_API_KEY)}")
    try:
        OPENAI_CLIENT = OpenAI(api_key=OPENAI_API_KEY)
        logger.info("✓ OpenAI client initialized successfully")
    except Exception as e:
        logger.warning(f"Failed to initialize OpenAI client: {e}")
        OPENAI_CLIENT = None

# Initialize Groq client if available
if USING_GROQ:
    logger.info(f"GROQ_API_KEY present: {bool(GROQ_API_KEY)}")
    try:
        GROQ_CLIENT = OpenAI(
            api_key=GROQ_API_KEY,
            base_url="https://api.groq.com/openai/v1"
        )
        logger.info("✓ Groq client initialized successfully")
    except Exception as e:
        logger.warning(f"Failed to initialize Groq client: {e}")
        GROQ_CLIENT = None

# Initialize Anthropic client if available (required for visualization)
if USING_ANTHROPIC:
    logger.info(f"ANTHROPIC_API_KEY present: {bool(ANTHROPIC_API_KEY)}")
    try:
        ANTHROPIC_CLIENT = Anthropic(api_key=ANTHROPIC_API_KEY)
        logger.info("✓ Anthropic client initialized successfully")
    except Exception as e:
        logger.warning(f"Failed to initialize Anthropic client: {e}")
        ANTHROPIC_CLIENT = None

# Set primary LLM (Priority: OpenAI > Groq - Anthropic is visualization-only)
if OPENAI_CLIENT:
    LLM_CLIENT = OPENAI_CLIENT
    PRIMARY_MODEL = OPENAI_CONTENT_MODEL
    LLM_SERVICE = "OpenAI"
    logger.info(f"✓ Using OpenAI ({PRIMARY_MODEL}) as primary LLM")
elif GROQ_CLIENT:
    LLM_CLIENT = GROQ_CLIENT
    PRIMARY_MODEL = GROQ_CONTENT_MODEL
    LLM_SERVICE = "Groq"
    logger.info(f"✓ Using Groq ({PRIMARY_MODEL}) as primary LLM")
else:
    LLM_CLIENT = None
    PRIMARY_MODEL = None
    LLM_SERVICE = None
    logger.warning("No content generation LLM available (OpenAI/Groq) - Anthropic is visualization-only")



# ============================================================================
# HELPER FUNCTIONS - BACKWARD COMPATIBILITY
# ============================================================================

def get_llm_client() -> Optional[OpenAI]:
    """Get the LLM client."""
    return LLM_CLIENT

def get_primary_model() -> str:
    """Get the primary model for reasoning tasks."""
    return PRIMARY_MODEL

def get_backup_model() -> str:
    """Get the backup model for reasoning tasks."""
    return PRIMARY_MODEL  # No backup, return primary

def get_llm_service() -> str:
    """Get the name of the LLM service being used."""
    return LLM_SERVICE

def is_using_openrouter() -> bool:
    """Check if OpenRouter is being used."""
    return False  # OpenRouter removed

def get_client_for_task(task_name: str):
    """
    Get the appropriate LLM client for a specific task.
    Task preferences:
      - content_generation -> prefer OpenAI, then Groq
      - visualization -> Anthropic only
    Returns: Client object (OpenAI-compatible or Anthropic)
    """
    t = (task_name or "").lower()
    if t == "visualization":
        return ANTHROPIC_CLIENT

    # Default: content generation - prefer OpenAI, then Groq (never Anthropic)
    if OPENAI_CLIENT:
        return OPENAI_CLIENT
    if GROQ_CLIENT:
        return GROQ_CLIENT
    return None  # No content generation client available

def get_model_for_task(task_name: str) -> str:
    """
    Get the appropriate model for a specific task.
    Strategy:
      - content_generation: prefer OpenAI, then Groq
      - visualization: Anthropic only
    """
    t = (task_name or "").lower()
    if t == "visualization":
        return ANTHROPIC_VISUALIZATION_MODEL

    # Default: content generation - prefer OpenAI, then Groq
    if USING_OPENAI:
        return OPENAI_CONTENT_MODEL
    if USING_GROQ:
        return GROQ_CONTENT_MODEL
    return OPENAI_CONTENT_MODEL or GROQ_CONTENT_MODEL

def get_service_for_task(task_name: str) -> str:
    """
    Get the service name for a specific task.
    """
    t = (task_name or "").lower()
    if t == "visualization":
        return "Anthropic"

    # Content generation: prefer OpenAI, then Groq
    if USING_OPENAI:
        return "OpenAI"
    if USING_GROQ:
        return "Groq"
    return "Unknown"

# ============================================================================
# AUDIO & NARRATION CONFIGURATION
# ============================================================================

AUDIO_PRIMARY_MODEL = "gemini-2.5-pro-preview-tts"
AUDIO_BACKUP_SERVICE = "regional_tts"
AUDIO_SERVICE = "Gemini"

NARRATION_PRIMARY_MODEL = PRIMARY_MODEL
NARRATION_BACKUP_MODEL = PRIMARY_MODEL  # No backup, use primary

def get_audio_model() -> str:
    """Get the primary audio generation model."""
    return AUDIO_PRIMARY_MODEL

def get_narration_model() -> str:
    """Get the model for generating narration text."""
    return NARRATION_PRIMARY_MODEL

def get_narration_backup_model() -> str:
    """Get the backup model for narration generation."""
    return NARRATION_BACKUP_MODEL

# ============================================================================
# PERSONALIZATION CONFIGURATION
# ============================================================================

PERSONALIZATION_CONFIG = {
    "audio_generation": {
        "primary_service": "Gemini",
        "primary_model": AUDIO_PRIMARY_MODEL,
        "backup_service": "regional_tts",
        "supports_voice_customization": True,
    },
    "reasoning": {
        "service": LLM_SERVICE,
        "primary_model": PRIMARY_MODEL,
        "backup_model": PRIMARY_MODEL,  # No backup
        "backup_service": LLM_SERVICE,  # No backup
        "supports_personalization": True,
        "supports_context_awareness": True,
    },
    "narration": {
        "service": LLM_SERVICE,
        "model": NARRATION_PRIMARY_MODEL,
        "backup_model": NARRATION_BACKUP_MODEL,
        "max_tokens": 500,
        "temperature": 0.6,
    },
}

# ============================================================================
# LOGGING
# ============================================================================

def log_model_configuration():
    """Log the current model configuration."""
    logger.info("=" * 70)
    logger.info("MODEL CONFIGURATION")
    logger.info("=" * 70)
    logger.info(f"Primary: {LLM_SERVICE} ({PRIMARY_MODEL})")
    logger.info(f"Visualization: Anthropic ({ANTHROPIC_VISUALIZATION_MODEL})")
    logger.info("=" * 70)