# utilties/llm_utils.py
"""
LLM utilities for Local Deep Research.

This module provides utility functions for working with language models
when the user's llm_config.py is missing or incomplete.
"""

import os
import logging
from typing import Dict, Any, Optional

# Setup logging
logger = logging.getLogger(__name__)

def get_model(
    model_name: Optional[str] = None,
    model_type: Optional[str] = None,
    temperature: Optional[float] = None,
    **kwargs
) -> Any:
    """
    Get a language model instance as fallback when llm_config.get_llm is not available.
    
    Args:
        model_name: Name of the model to use
        model_type: Type of the model provider
        temperature: Model temperature
        **kwargs: Additional parameters
        
    Returns:
        LangChain language model instance
    """
    # Get default values from kwargs or use reasonable defaults
    model_name = model_name or kwargs.get('DEFAULT_MODEL', 'mistral')
    model_type = model_type or kwargs.get('DEFAULT_MODEL_TYPE', 'ollama')
    temperature = temperature or kwargs.get('DEFAULT_TEMPERATURE', 0.7)
    max_tokens = kwargs.get('max_tokens', kwargs.get('MAX_TOKENS', 30000))
    
    # Common parameters
    common_params = {
        "temperature": temperature,
        "max_tokens": max_tokens,
    }
    
    # Add additional kwargs
    for key, value in kwargs.items():
        if key not in ['DEFAULT_MODEL', 'DEFAULT_MODEL_TYPE', 'DEFAULT_TEMPERATURE', 'MAX_TOKENS']:
            common_params[key] = value
    
    # Try to load the model based on type
    if model_type == "ollama":
        try:
            from langchain_ollama import ChatOllama
            return ChatOllama(model=model_name, **common_params)
        except ImportError:
            try:
                from langchain_community.llms import Ollama
                return Ollama(model=model_name, **common_params)
            except ImportError:
                logger.error("Neither langchain_ollama nor langchain_community.llms.Ollama available")
                raise
    
    elif model_type == "openai":
        try:
            from langchain_openai import ChatOpenAI
            api_key = os.getenv("OPENAI_API_KEY")
            if not api_key:
                raise ValueError("OPENAI_API_KEY environment variable not set")
            return ChatOpenAI(model=model_name, api_key=api_key, **common_params)
        except ImportError:
            logger.error("langchain_openai not available")
            raise
    
    elif model_type == "anthropic":
        try:
            from langchain_anthropic import ChatAnthropic
            api_key = os.getenv("ANTHROPIC_API_KEY")
            if not api_key:
                raise ValueError("ANTHROPIC_API_KEY environment variable not set")
            return ChatAnthropic(model=model_name, anthropic_api_key=api_key, **common_params)
        except ImportError:
            logger.error("langchain_anthropic not available")
            raise
    
    elif model_type == "openai_endpoint":
        try:
            from langchain_openai import ChatOpenAI
            api_key = os.getenv("OPENAI_ENDPOINT_API_KEY")
            if not api_key:
                raise ValueError("OPENAI_ENDPOINT_API_KEY environment variable not set")
            
            endpoint_url = kwargs.get("OPENAI_ENDPOINT_URL", "https://openrouter.ai/api/v1")
            
            if model_name is None and not kwargs.get("OPENAI_ENDPOINT_REQUIRES_MODEL", True):
                return ChatOpenAI(api_key=api_key, openai_api_base=endpoint_url, **common_params)
            else:
                return ChatOpenAI(model=model_name, api_key=api_key, openai_api_base=endpoint_url, **common_params)
        except ImportError:
            logger.error("langchain_openai not available")
            raise
    
    # Default fallback
    try:
        from langchain_ollama import ChatOllama
        logger.warning(f"Unknown model type '{model_type}', defaulting to Ollama")
        return ChatOllama(model=model_name, **common_params)
    except (ImportError, Exception) as e:
        logger.error(f"Failed to load any model: {e}")
        
        # Last resort: create a dummy model
        try:
            from langchain_community.llms.fake import FakeListLLM
            return FakeListLLM(responses=["No language models are available. Please install Ollama or set up API keys."])
        except ImportError:
            raise ValueError("No language models available and could not create dummy model")