"""
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

# Initialize environment
logger = logging.getLogger(__name__)

# Valid provider options
VALID_PROVIDERS = ["ollama", "openai", "anthropic", "vllm", "openai_endpoint", "lmstudio", "llamacpp", "none"]

# ================================
# 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 settings.llm.model)
        temperature: Model temperature (if None, uses settings.llm.temperature)
        provider: Provider to use (if None, uses settings.llm.provider)
    
    Returns:
        A LangChain LLM instance
    """
    # Use settings values for parameters if not provided
    if model_name is None:
        model_name = settings.llm.model
    
    if temperature is None:
        temperature = settings.llm.temperature
        
    if provider is None:
        provider = settings.llm.provider.lower()
        if provider not in VALID_PROVIDERS:
            logger.error(f"Invalid provider in settings: {provider}")
            raise ValueError(f"Invalid provider: {provider}. Must be one of: {VALID_PROVIDERS}")
    
    # Common parameters for all models
    common_params = {
        "temperature": temperature,
        "max_tokens": settings.llm.max_tokens,
    }
    
    # Handle different providers
    if provider == "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 == "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 == "openai_endpoint":
        api_key = settings.get('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)
        
        # Get endpoint URL from settings
        openai_endpoint_url = settings.llm.openai_endpoint_url
        
        return ChatOpenAI(
            model=model_name, 
            api_key=api_key,
            openai_api_base=openai_endpoint_url, 
            **common_params
        )
    
    elif provider == "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 == "ollama":
        try:
            # Use the configurable Ollama base URL
            base_url = settings.get('OLLAMA_BASE_URL', settings.llm.get('ollama_base_url', 'http://localhost:11434'))
            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)
    
    elif provider == "lmstudio":

            # LM Studio supports OpenAI API format, so we can use ChatOpenAI directly
            lmstudio_url = settings.llm.get('lmstudio_url', "http://localhost:1234")
            
            return ChatOpenAI(
                model=model_name,
                api_key="lm-studio",  # LM Studio doesn't require a real API key
                base_url=f"{lmstudio_url}/v1",  # Use the configured URL with /v1 endpoint
                temperature=temperature,
                max_tokens=settings.llm.max_tokens
            )

 
    elif provider == "llamacpp":

            # Import LlamaCpp
            from langchain_community.llms import LlamaCpp
            
            # Get LlamaCpp model path from settings
            model_path = settings.llm.get('llamacpp_model_path', "")
            if not model_path:
                logger.error("llamacpp_model_path not set in settings")
                raise ValueError("llamacpp_model_path not set in settings.toml")
                
            # Get additional LlamaCpp parameters
            n_gpu_layers = settings.llm.get('llamacpp_n_gpu_layers', 1)
            n_batch = settings.llm.get('llamacpp_n_batch', 512)
            f16_kv = settings.llm.get('llamacpp_f16_kv', True)
            
            # Create LlamaCpp instance
            return LlamaCpp(
                model_path=model_path,
                temperature=temperature,
                max_tokens=settings.llm.max_tokens,
                n_gpu_layers=n_gpu_layers,
                n_batch=n_batch,
                f16_kv=f16_kv,
                verbose=True
            )
    
    else:
        return get_fallback_model(temperature)

def get_fallback_model(temperature=None):
    """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"
    
    if is_lmstudio_available():
        providers["lmstudio"] = "LM Studio (local models)"
        
    if is_llamacpp_available():
        providers["llamacpp"] = "LlamaCpp (local models)"
    
    # 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.get('OPENAI_API_KEY', '')
        return bool(api_key)
    except:
        return False

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

def is_openai_endpoint_available():
    """Check if OpenAI endpoint is available"""
    try:
        api_key = settings.get('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', settings.llm.get('ollama_base_url', 'http://localhost:11434'))
        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 is_lmstudio_available():
    """Check if LM Studio is available"""
    try:
        import requests
        lmstudio_url = settings.llm.get('lmstudio_url', 'http://localhost:1234')
        # LM Studio typically uses OpenAI-compatible endpoints
        response = requests.get(f"{lmstudio_url}/v1/models", timeout=1.0)
        return response.status_code == 200
    except:
        return False

def is_llamacpp_available():
    """Check if LlamaCpp is available and configured"""
    try:
        from langchain_community.llms import LlamaCpp
        model_path = settings.llm.get('llamacpp_model_path', '')
        return bool(model_path) and os.path.exists(model_path)
    except:
        return False

def get_available_providers():
    """Get dictionary of available providers"""
    return get_available_provider_types()

# Log which providers are available
AVAILABLE_PROVIDERS = get_available_providers()
logger.info(f"Available providers: {list(AVAILABLE_PROVIDERS.keys())}")

# Check if selected provider is available
selected_provider = settings.llm.provider.lower()
if selected_provider not in AVAILABLE_PROVIDERS and selected_provider != "none":
    logger.warning(f"Selected provider {selected_provider} is not available.")