"""
LLM configuration for Local Deep Research.

This file controls which language models are available and how they're configured.
You can customize model selection, parameters, and fallbacks here.
"""

from langchain_anthropic import ChatAnthropic
from langchain_openai import ChatOpenAI
from langchain_ollama import ChatOllama
from langchain_community.llms import VLLM
from local_deep_research.config import settings
import os
import logging
from enum import Enum, auto

# Initialize environment
logger = logging.getLogger(__name__)

# Provider enum
class ModelProvider(Enum):
    OLLAMA = auto()
    OPENAI = auto()
    ANTHROPIC = auto()
    VLLM = auto()
    OPENAI_ENDPOINT = auto()
    NONE = auto()

# ================================
# USER CONFIGURATION SECTION
# ================================

# Set your preferred model provider here
DEFAULT_PROVIDER = ModelProvider.OLLAMA  # Change this to your preferred provider

# Set your default model name here
DEFAULT_MODEL = "mistral"  # Your default model

# Set default model parameters
DEFAULT_TEMPERATURE = 0.7
MAX_TOKENS = 30000

# Server URLs
OPENAI_ENDPOINT_URL = "https://openrouter.ai/api/v1"  # For OpenRouter or compatible services
OLLAMA_BASE_URL = "http://localhost:11434"  # URL for Ollama server




# ================================
# LLM FUNCTIONS
# ================================



    

def get_llm(model_name=None, temperature=None, provider=None):
    """
    Get LLM instance based on model name and provider.
    
    Args:
        model_name: Name of the model to use (if None, uses DEFAULT_MODEL)
        temperature: Model temperature (if None, uses DEFAULT_TEMPERATURE)
        provider: Provider to use (if None, uses DEFAULT_PROVIDER)
    
    Returns:
        A LangChain LLM instance
    """
    if model_name is None:
        model_name = DEFAULT_MODEL
    
    if temperature is None:
        temperature = DEFAULT_TEMPERATURE
        
    if provider is None:
        provider = DEFAULT_PROVIDER
    
    # Common parameters for all models
    common_params = {
        "temperature": temperature,
        "max_tokens": MAX_TOKENS,
    }
    
    # Handle different providers
    if provider == ModelProvider.ANTHROPIC:
        api_key = settings.get('ANTHROPIC_API_KEY', '')
        if not api_key:
            logger.warning("ANTHROPIC_API_KEY not found. Falling back to default model.")
            return get_fallback_model(temperature)
        
        return ChatAnthropic(
            model=model_name, anthropic_api_key=api_key, **common_params
        )
    
    elif provider == ModelProvider.OPENAI:
        api_key = settings.get('OPENAI_API_KEY', '')
        if not api_key:
            logger.warning("OPENAI_API_KEY not found. Falling back to default model.")
            return get_fallback_model(temperature)
        
        return ChatOpenAI(model=model_name, api_key=api_key, **common_params)
    
    elif provider == ModelProvider.OPENAI_ENDPOINT:
        api_key = settings.OPENAI_ENDPOINT_API_KEY

        if not api_key:
            logger.warning("OPENAI_ENDPOINT_API_KEY not found. Falling back to default model.")
            return get_fallback_model(temperature)
        
        return ChatOpenAI(
            model=model_name, 
            api_key=api_key,
            openai_api_base=OPENAI_ENDPOINT_URL, 
            **common_params
        )
    
    elif provider == ModelProvider.VLLM:
        try:
            return VLLM(
                model=model_name,
                trust_remote_code=True,
                max_new_tokens=128,
                top_k=10,
                top_p=0.95,
                temperature=temperature,
            )
        except Exception as e:
            logger.error(f"Error loading VLLM model: {e}")
            logger.warning("Falling back.")
            return get_fallback_model(temperature)
    
    elif provider == ModelProvider.OLLAMA:
        try:
            # Use the configurable Ollama base URL
            base_url = settings.get('OLLAMA_BASE_URL', OLLAMA_BASE_URL)
            return ChatOllama(model=model_name, base_url=base_url, **common_params)
        except Exception as e:
            logger.error(f"Error loading Ollama model: {e}")
            return get_fallback_model(temperature)
    
    else:
        return get_fallback_model(temperature)

def get_fallback_model(temperature=DEFAULT_TEMPERATURE):
    """Create a dummy model for when no providers are available"""
    from langchain_community.llms.fake import FakeListLLM
    return FakeListLLM(
        responses=["No language models are available. Please install Ollama or set up API keys."]
    )

# ================================
# COMPATIBILITY FUNCTIONS
# ================================

def get_available_provider_types():
    """Return available model providers"""
    providers = {}
    
    if is_ollama_available():
        providers["ollama"] = "Ollama (local models)"
    
    if is_openai_available():
        providers["openai"] = "OpenAI API"
    
    if is_anthropic_available():
        providers["anthropic"] = "Anthropic API"
    
    if is_openai_endpoint_available():
        providers["openai_endpoint"] = "OpenAI-compatible Endpoint"
    
    # Check for VLLM capability
    try:
        import torch
        import transformers
        providers["vllm"] = "VLLM (local models)"
    except ImportError:
        pass
    
    # Default fallback
    if not providers:
        providers["none"] = "No model providers available"
    
    return providers




# ================================
# HELPER FUNCTIONS
# ================================

def is_openai_available():
    """Check if OpenAI is available"""
    try:
        api_key = settings.api_keys.get('OPENAI_API_KEY', '')
        return bool(api_key)
    except:
        return False

def is_anthropic_available():
    """Check if Anthropic is available"""
    try:
        api_key = settings.api_keys.get('ANTHROPIC_API_KEY', '')
        return bool(api_key)
    except:
        return False



def is_openai_endpoint_available():
    """Check if OpenAI endpoint is available"""
    print(os.getenv("OPENAI_ENDPOINT_API_KEY"))
    try:
        api_key = settings.OPENAI_ENDPOINT_API_KEY
        return bool(api_key) 
    except:
        return False

def is_ollama_available():
    """Check if Ollama is running"""
    try:
        import requests
        base_url = settings.get('OLLAMA_BASE_URL', OLLAMA_BASE_URL)
        response = requests.get(f"{base_url}/api/tags", timeout=1.0)
        return response.status_code == 200
    except:
        return False

def is_vllm_available():
    """Check if VLLM capability is available"""
    try:
        import torch
        import transformers
        return True
    except ImportError:
        return False

def get_available_providers():
    """Get dictionary of available providers"""
    providers = {}
    
    if is_ollama_available():
        providers[ModelProvider.OLLAMA] = "Ollama (local models)"
    
    if is_openai_available():
        providers[ModelProvider.OPENAI] = "OpenAI API"
    
    if is_anthropic_available():
        providers[ModelProvider.ANTHROPIC] = "Anthropic API"
    
    if is_openai_endpoint_available():
        providers[ModelProvider.OPENAI_ENDPOINT] = "OpenAI-compatible Endpoint"
    
    if is_vllm_available():
        providers[ModelProvider.VLLM] = "VLLM (local models)"
    
    if not providers:
        providers[ModelProvider.NONE] = "No model providers available"
    
    return providers

# Log which providers are available
AVAILABLE_PROVIDERS = get_available_providers()
logger.info(f"Available providers: {[p.name for p in AVAILABLE_PROVIDERS.keys()]}")

# Check if selected provider is available
if DEFAULT_PROVIDER not in AVAILABLE_PROVIDERS and DEFAULT_PROVIDER != ModelProvider.NONE:
    logger.warning(f"Selected provider {DEFAULT_PROVIDER.name} is not available.")
