#!/usr/bin/env python3
"""
LLM Configuration Module

Loads LLM configuration from .env file or environment variables.
Supports multiple providers via LiteLLM.
"""

import os
from pathlib import Path
from typing import Dict, Optional, Any
from dotenv import load_dotenv

PROJECT_ROOT = Path(__file__).resolve().parents[2]
ROOT_ENV_FILE = PROJECT_ROOT / ".env"

if ROOT_ENV_FILE.exists():
    load_dotenv(ROOT_ENV_FILE)

# Allowed LLM providers
ALLOWED_LLM_PROVIDERS = {
    "openai", "azure", "anthropic", "mistral", "codestral",
    "groq", "openrouter", "huggingface", "cohere", "bedrock",
    "vertex_ai", "gemini", "ollama", "minimax", "openai-compatible",
    "openai_compatible"
}


def get_model_name(provider: Optional[str], model: Optional[str]) -> str:
    """
    Construct the model name in LiteLLM format.
    
    Args:
        provider: Provider name (e.g., "openai", "azure", "anthropic")
        model: Model name (e.g., "gpt-4o", "claude-3-opus", "openrouter/google/gemini-pro")
    
    Returns:
        Model name in LiteLLM format (e.g., "gpt-4o" or "azure/gpt-4o")
    """
    if not model:
        return "gpt-4o"  # Default fallback
    
    # For OpenAI, return as-is (no prefix needed)
    if provider == "openai":
        return model
    
    # For Azure, ensure model looks like "azure/<deployment_name>"
    if provider == "azure":
        if model.startswith("azure/"):
            return model
        return f"azure/{model}"
    
    # For all other providers, add provider/ prefix if not already present
    if provider:
        if model.startswith(f"{provider}/"):
            return model  # Already has correct prefix (e.g., openrouter/minimax/...)
        return f"{provider}/{model}"
    
    return model


def is_openai_official_model(model: str) -> bool:
    """
    Return True for model names LiteLLM can already resolve as OpenAI models.
    """
    official_prefixes = (
        "gpt-",
        "o1",
        "o3",
        "o4",
        "chatgpt-",
        "text-",
        "davinci",
        "babbage",
        "curie",
        "ada",
    )
    return model.startswith(official_prefixes)


def get_openai_model_name(model: Optional[str], endpoint: Optional[str]) -> str:
    """
    Format OpenAI and OpenAI-compatible model names for LiteLLM.

    Custom OpenAI-compatible endpoints often use model names like "qwen-plus".
    LiteLLM needs these explicitly routed through the OpenAI adapter as
    "openai/qwen-plus"; official OpenAI models can remain unchanged.
    """
    model_name = model or "gpt-4o"
    if model_name.startswith("openai/"):
        return model_name
    if endpoint and not is_openai_official_model(model_name):
        return f"openai/{model_name}"
    return model_name


def load_llm_config() -> Dict[str, Any]:
    """
    Load LLM configuration from .env file or environment variables.
    
    Returns:
        Dictionary with LLM configuration:
        {
            "provider": str,
            "model": str,
            "api_key": str,
            "endpoint": Optional[str],
            "api_version": Optional[str],
            "temperature": float,
            "top_p": float
        }
    
    Raises:
        ValueError: If required configuration is missing
    """
    # Determine provider
    provider = os.getenv("PROVIDER", "openai").lower()
    
    # Normalize aliases to canonical provider name
    if provider == "google":
        provider = "gemini"
    elif provider in {"openai-compatible", "openai_compatible"}:
        provider = "openai"
    
    # Validate provider is in allowed list
    if provider not in ALLOWED_LLM_PROVIDERS:
        raise ValueError(
            f"Provider '{provider}' is not supported. "
            f"Allowed providers: {', '.join(sorted(ALLOWED_LLM_PROVIDERS))}"
        )
    
    # Get model name
    model = os.getenv("MODEL", "gpt-4o")
    
    # Get API key and provider-specific config based on provider
    api_key = None
    endpoint = None
    api_version = None
    
    if provider == "openai":
        api_key = os.getenv("OPENAI_API_KEY")
        endpoint = (
            os.getenv("OPENAI_BASE_URL")
            or os.getenv("OPENAI_API_BASE")
            or os.getenv("OPENAI_COMPATIBLE_BASE_URL")
        )
    
    elif provider == "azure":
        api_key = os.getenv("AZURE_OPENAI_API_KEY") or os.getenv("AZURE_API_KEY")
        endpoint = os.getenv("AZURE_OPENAI_ENDPOINT") or os.getenv("AZURE_API_BASE")
        api_version = os.getenv("AZURE_OPENAI_API_VERSION") or os.getenv("AZURE_API_VERSION", "2024-08-01-preview")
    
    elif provider == "anthropic":
        api_key = os.getenv("ANTHROPIC_API_KEY")
    
    elif provider == "gemini":
        api_key = os.getenv("GOOGLE_API_KEY")
    
    elif provider == "mistral":
        api_key = os.getenv("MISTRAL_API_KEY")
    
    elif provider == "codestral":
        # Codestral uses Mistral API key
        api_key = os.getenv("MISTRAL_API_KEY")
    
    elif provider == "groq":
        api_key = os.getenv("GROQ_API_KEY")
    
    elif provider == "openrouter":
        # OpenRouter uses the OpenAI-compatible API format.
        # Use OpenAI provider with OpenRouter base URL to bypass LiteLLM's
        # buggy openrouter adapter (which fails to send auth headers).
        api_key = os.getenv("OPENROUTER_API_KEY")
        endpoint = "https://openrouter.ai/api/v1"
        # Force OpenAI as the effective provider so get_model_name() passes the
        # raw model name through (openrouter model IDs like "minimax/minimax-m2.7"
        # are already in the correct format and must NOT be prefixed again).
        provider = "openai"

    elif provider == "minimax":
        # LiteLLM has a native MiniMax provider. Route through it directly
        # instead of pretending MiniMax is "openai", which breaks provider
        # resolution for newer MiniMax model IDs like "MiniMax-M2.7".
        #
        # Accept both LiteLLM-style env vars and MiniMax's OpenAI-compatible
        # docs-style env vars as fallbacks.
        api_key = os.getenv("MINIMAX_API_KEY") or os.getenv("OPENAI_API_KEY")
        endpoint = (
            os.getenv("MINIMAX_API_BASE")
            or os.getenv("OPENAI_BASE_URL")
            or os.getenv("OPENAI_API_BASE")
            or "https://api.minimaxi.com/v1"
        )
    
    elif provider == "huggingface":
        api_key = os.getenv("HUGGINGFACE_API_KEY")
    
    elif provider == "cohere":
        api_key = os.getenv("COHERE_API_KEY") or os.getenv("CO_API_KEY")
    
    elif provider == "bedrock":
        # Bedrock supports multiple authentication methods:
        # 1. AWS SSO via profile (recommended): AWS_PROFILE
        # 2. Static credentials: AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY
        # 3. Temporary credentials: Above + AWS_SESSION_TOKEN
        aws_profile = os.getenv("AWS_PROFILE")
        aws_region = os.getenv("AWS_REGION_NAME", "us-east-1")
        
        if aws_profile:
            # Profile-based auth (SSO, IAM roles) - SDK handles credential resolution
            api_key = "bedrock_profile_auth"
        else:
            # Static or temporary credentials
            api_key = os.getenv("AWS_ACCESS_KEY_ID")
        
        # Store region in endpoint field for Bedrock
        endpoint = aws_region
    
    elif provider == "vertex_ai":
        # Vertex AI uses GCP credentials (service account JSON or GOOGLE_APPLICATION_CREDENTIALS)
        # No API key needed, but we set a placeholder to pass validation
        gcp_creds = os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
        if not gcp_creds and not Path("~/.config/gcloud/application_default_credentials.json").expanduser().exists():
            raise ValueError(
                "GCP credentials not found. Set GOOGLE_APPLICATION_CREDENTIALS or run 'gcloud auth application-default login'"
            )
        api_key = "vertex_ai_placeholder"
    
    elif provider == "ollama":
        # Ollama uses OLLAMA_BASE_URL (defaults to http://localhost:11434)
        endpoint = os.getenv("OLLAMA_BASE_URL", "http://localhost:11434")
        # Ollama doesn't require API key, but we set a placeholder to pass validation
        api_key = "ollama_placeholder"
    
    # Get optional parameters
    temperature = float(os.getenv("LLM_TEMPERATURE", "0.2"))
    top_p = float(os.getenv("LLM_TOP_P", "0.2"))
    timeout = int(os.getenv("LLM_TIMEOUT", "300"))
    max_retries = int(os.getenv("LLM_MAX_RETRIES", "3"))

    config = {
        "provider": provider,
        "model": (
            get_openai_model_name(model, endpoint)
            if provider == "openai"
            else get_model_name(provider, model)
        ),
        "api_key": api_key,
        "temperature": temperature,
        "top_p": top_p,
        "timeout": timeout,
        "max_retries": max_retries,
        "api_base": endpoint if endpoint else None
    }
    
    # Add provider-specific fields
    if endpoint:
        config["endpoint"] = endpoint
    if api_version:
        config["api_version"] = api_version
    
    # Special handling for Bedrock (store AWS credentials and profile)
    if provider == "bedrock":
        config["aws_region"] = endpoint  # Store region in endpoint field
        
        # Profile-based authentication (AWS SSO, IAM roles)
        aws_profile = os.getenv("AWS_PROFILE")
        if aws_profile:
            config["aws_profile"] = aws_profile
        else:
            # Static or temporary credentials
            config["aws_secret_access_key"] = os.getenv("AWS_SECRET_ACCESS_KEY")
            aws_session_token = os.getenv("AWS_SESSION_TOKEN")
            if aws_session_token:
                config["aws_session_token"] = aws_session_token
    
    # Special handling for Vertex AI (store GCP project/location if provided)
    if provider == "vertex_ai":
        if os.getenv("GCP_PROJECT_ID"):
            config["gcp_project_id"] = os.getenv("GCP_PROJECT_ID")
        if os.getenv("GCP_LOCATION"):
            config["gcp_location"] = os.getenv("GCP_LOCATION")
    
    return config


