"""
LLM interface module using LiteLLM.

This module provides a unified interface to various LLM providers through LiteLLM,
enabling query enhancement, classification, and other LLM-powered functionality.
"""

import os
import json
from typing import Dict, Any, List, Optional, Tuple, Union

import litellm
from litellm import completion

from config.config import get_config


class LLMInterface:
    """Interface for interacting with LLMs through LiteLLM."""
    
    def __init__(self, model_name: Optional[str] = None):
        """
        Initialize the LLM interface.
        
        Args:
            model_name: Name of the LLM model to use. If None, uses the default model
                       from configuration.
        """
        self.config = get_config()
        
        # Use specified model or default from config
        self.model_name = model_name or self.config.config_data.get('default_model', 'gpt-3.5-turbo')
        
        # Get model-specific configuration
        self.model_config = self.config.get_model_config(self.model_name)
        
        # Set up LiteLLM with the appropriate provider
        self._setup_provider()
    
    def _setup_provider(self) -> None:
        """Set up the LLM provider based on the model configuration."""
        provider = self.model_config.get('provider', 'openai')
        
        try:
            # Get API key for the provider
            api_key = self.config.get_api_key(provider)
            
            # Set environment variable for the provider
            if provider.lower() == 'google':
                os.environ["GEMINI_API_KEY"] = api_key
            else:
                os.environ[f"{provider.upper()}_API_KEY"] = api_key
            
            print(f"LLM interface initialized with model: {self.model_name} (provider: {provider})")
        except ValueError as e:
            print(f"Error setting up LLM provider: {e}")
    
    def _get_completion_params(self) -> Dict[str, Any]:
        """
        Get parameters for LLM completion based on model configuration.
        
        Returns:
            Dictionary of parameters for LiteLLM completion
        """
        params = {
            'temperature': self.model_config.get('temperature', 0.7),
            'max_tokens': self.model_config.get('max_tokens', 1000),
            'top_p': self.model_config.get('top_p', 1.0)
        }
        
        # Handle different provider configurations
        provider = self.model_config.get('provider', 'openai')
        
        if provider == 'azure':
            # Azure OpenAI requires special handling
            deployment_name = self.model_config.get('deployment_name')
            api_version = self.model_config.get('api_version')
            endpoint = self.model_config.get('endpoint')
            
            if deployment_name and endpoint:
                # Format: azure/deployment_name
                params['model'] = f"azure/{deployment_name}"
                
                # Set Azure-specific environment variables if not already set
                if 'AZURE_API_BASE' not in os.environ and endpoint:
                    os.environ['AZURE_API_BASE'] = endpoint
                
                if 'AZURE_API_VERSION' not in os.environ and api_version:
                    os.environ['AZURE_API_VERSION'] = api_version
            else:
                # Fall back to default model if Azure config is incomplete
                params['model'] = self.model_name
        elif provider in ['ollama', 'groq', 'openrouter'] or self.model_config.get('endpoint'):
            # For providers with custom endpoints
            params['model'] = self.model_config.get('model_name', self.model_name)
            params['api_base'] = self.model_config.get('endpoint')
            
            # Special handling for OpenRouter
            if provider == 'openrouter':
                # Set HTTP headers for OpenRouter if needed
                params['headers'] = {
                    'HTTP-Referer': 'https://sim-search.app',  # Replace with your actual app URL
                    'X-Title': 'Intelligent Research System'   # Replace with your actual app name
                }
        elif provider == 'google':
            # Special handling for Google Gemini models
            params['model'] = f"gemini/{self.model_config.get('model_name', self.model_name)}"
            # Google Gemini uses a different API base
            params['api_base'] = self.model_config.get('endpoint', 'https://generativelanguage.googleapis.com/v1')
        else:
            # Standard provider (OpenAI, Anthropic, etc.)
            params['model'] = self.model_name
        
        return params
    
    def generate_completion(self, messages: List[Dict[str, str]], stream: bool = False) -> Union[str, Any]:
        """
        Generate a completion using the configured LLM.
        
        Args:
            messages: List of message dictionaries with 'role' and 'content' keys
            stream: Whether to stream the response
            
        Returns:
            If stream is False, returns the completion text as a string
            If stream is True, returns the completion response object for streaming
        """
        try:
            params = self._get_completion_params()
            params['messages'] = messages
            params['stream'] = stream
            
            response = completion(**params)
            
            if stream:
                return response
            else:
                return response.choices[0].message.content
        except Exception as e:
            print(f"Error generating completion: {e}")
            return f"Error: {str(e)}"
    
    def enhance_query(self, query: str) -> str:
        """
        Enhance a user query using the LLM.
        
        Args:
            query: The raw user query
            
        Returns:
            Enhanced query with additional context and structure
        """
        # Get the model assigned to this specific function
        model_name = self.config.get_module_model('query_processing', 'enhance_query')
        
        # Create a new interface with the assigned model if different from current
        if model_name != self.model_name:
            interface = LLMInterface(model_name)
            return interface._enhance_query_impl(query)
        
        return self._enhance_query_impl(query)
    
    def _enhance_query_impl(self, query: str) -> str:
        """Implementation of query enhancement."""
        messages = [
            {"role": "system", "content": "You are an AI research assistant. Your task is to enhance the user's query by adding relevant context, clarifying ambiguities, and expanding key terms. Maintain the original intent of the query while making it more comprehensive and precise. Return ONLY the enhanced query text without any explanations, introductions, or additional text. The enhanced query should be ready to be sent directly to a search engine."},
            {"role": "user", "content": f"Enhance this research query: {query}"}
        ]
        
        return self.generate_completion(messages)
    
    def classify_query(self, query: str) -> Dict[str, Any]:
        """
        Classify a query to determine its type, intent, and key entities.
        
        Args:
            query: The user query to classify
            
        Returns:
            Dictionary containing query classification information
        """
        # Get the model assigned to this specific function
        model_name = self.config.get_module_model('query_processing', 'classify_query')
        
        # Create a new interface with the assigned model if different from current
        if model_name != self.model_name:
            interface = LLMInterface(model_name)
            return interface._classify_query_impl(query)
        
        return self._classify_query_impl(query)
    
    def _classify_query_impl(self, query: str) -> Dict[str, Any]:
        """Implementation of query classification."""
        messages = [
            {"role": "system", "content": "You are an AI research assistant. Analyze the user's query and classify it according to type (factual, exploratory, comparative, etc.), intent, and key entities. Respond with a JSON object containing these classifications."},
            {"role": "user", "content": f"Classify this research query: {query}"}
        ]
        
        response = self.generate_completion(messages)
        
        try:
            # Try to parse as JSON
            classification = json.loads(response)
            return classification
        except json.JSONDecodeError:
            # If not valid JSON, return a basic classification
            return {
                "type": "unknown",
                "intent": "research",
                "entities": [query],
                "error": "Failed to parse LLM response as JSON"
            }
    
    def generate_search_queries(self, query: str, search_engines: List[str]) -> Dict[str, List[str]]:
        """
        Generate optimized search queries for different search engines.
        
        Args:
            query: The original user query
            search_engines: List of search engines to generate queries for
            
        Returns:
            Dictionary mapping search engines to lists of optimized queries
        """
        # Get the model assigned to this specific function
        model_name = self.config.get_module_model('query_processing', 'generate_search_queries')
        
        # Create a new interface with the assigned model if different from current
        if model_name != self.model_name:
            interface = LLMInterface(model_name)
            return interface._generate_search_queries_impl(query, search_engines)
        
        return self._generate_search_queries_impl(query, search_engines)
    
    def _generate_search_queries_impl(self, query: str, search_engines: List[str]) -> Dict[str, List[str]]:
        """Implementation of search query generation."""
        engines_str = ", ".join(search_engines)
        
        messages = [
            {"role": "system", "content": f"You are an AI research assistant. Generate optimized search queries for the following search engines: {engines_str}. For each search engine, provide 3 variations of the query that are optimized for that engine's search algorithm and will yield comprehensive results."},
            {"role": "user", "content": f"Generate optimized search queries for this research topic: {query}"}
        ]
        
        response = self.generate_completion(messages)
        
        try:
            # Try to parse as JSON
            queries = json.loads(response)
            return queries
        except json.JSONDecodeError:
            # If not valid JSON, return a basic query set
            return {engine: [query] for engine in search_engines}


# Create a singleton instance for global use
llm_interface = LLMInterface()


def get_llm_interface(model_name: Optional[str] = None) -> LLMInterface:
    """
    Get the global LLM interface instance or create a new one with a specific model.
    
    Args:
        model_name: Optional model name to use instead of the default
        
    Returns:
        LLMInterface instance
    """
    if model_name:
        return LLMInterface(model_name)
    return llm_interface