import wikipedia
from typing import Dict, List, Any, Optional
from langchain_core.language_models import BaseLLM
import logging

from local_deep_research.web_search_engines.search_engine_base import BaseSearchEngine
from local_deep_research import config

# Setup logging
logger = logging.getLogger(__name__)


class WikipediaSearchEngine(BaseSearchEngine):
    """Wikipedia search engine implementation with two-phase approach"""
    
    def __init__(self, 
                max_results: int = 10, 
                language: str = "en", 
                include_content: bool = True,
                sentences: int = 5,
                llm: Optional[BaseLLM] = None,
                max_filtered_results: Optional[int] = None,
                **kwargs):
        """
        Initialize the Wikipedia search engine.
        
        Args:
            max_results: Maximum number of search results
            language: Language code for Wikipedia (e.g., 'en', 'fr', 'es')
            include_content: Whether to include full page content in results
            sentences: Number of sentences to include in summary
            llm: Language model for relevance filtering
            max_filtered_results: Maximum number of results to keep after filtering
            **kwargs: Additional parameters (ignored but accepted for compatibility)
        """
        # Initialize the BaseSearchEngine with LLM, max_filtered_results, and max_results
        super().__init__(llm=llm, max_filtered_results=max_filtered_results, max_results=max_results)
        self.include_content = include_content
        self.sentences = sentences
        
        # Set the Wikipedia language
        wikipedia.set_lang(language)
    
    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information (titles and summaries) for Wikipedia pages.
        
        Args:
            query: The search query
            
        Returns:
            List of preview dictionaries
        """
        logger.info(f"Getting Wikipedia page previews for query: {query}")
        
        try:
            # Get search results (just titles)
            search_results = wikipedia.search(query, results=self.max_results)
            
            logger.info(f"Found {len(search_results)} Wikipedia results: {search_results}")
            
            if not search_results:
                logger.info(f"No Wikipedia results found for query: {query}")
                return []
            
            # Create a cache for full pages (will be populated on-demand)
            self._page_cache = {}
            
            # Generate previews with summaries
            previews = []
            for title in search_results:
                try:
                    # Get just the summary, with auto_suggest=False to be more precise
                    summary = None
                    try:
                        summary = wikipedia.summary(title, sentences=self.sentences, auto_suggest=False)
                    except wikipedia.exceptions.DisambiguationError as e:
                        # If disambiguation error, try the first option
                        if e.options and len(e.options) > 0:
                            logger.info(f"Disambiguation for '{title}', trying first option: {e.options[0]}")
                            try:
                                summary = wikipedia.summary(e.options[0], sentences=self.sentences, auto_suggest=False)
                                title = e.options[0]  # Use the new title
                            except Exception as inner_e:
                                logger.error(f"Error with disambiguation option: {inner_e}")
                                continue
                        else:
                            logger.warning(f"Disambiguation with no options for '{title}'")
                            continue
                    
                    if summary:
                        preview = {
                            "id": title,  # Use title as ID
                            "title": title,
                            "snippet": summary,
                            "link": f"https://en.wikipedia.org/wiki/{title.replace(' ', '_')}"
                        }
                        
                        previews.append(preview)
                    
                except (wikipedia.exceptions.PageError,
                       wikipedia.exceptions.WikipediaException) as e:
                    # Skip pages with errors
                    logger.warning(f"Error getting summary for '{title}': {e}")
                    continue
                except Exception as e:
                    logger.error(f"Unexpected error for '{title}': {e}")
                    continue
            
            logger.info(f"Successfully created {len(previews)} previews from Wikipedia")
            return previews
            
        except Exception as e:
            logger.error(f"Error getting Wikipedia previews: {e}")
            return []
    
    def _get_full_content(self, relevant_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant Wikipedia pages.
        
        Args:
            relevant_items: List of relevant preview dictionaries
            
        Returns:
            List of result dictionaries with full content
        """
        # Check if we should add full content
        if hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY:
            logger.info("Snippet-only mode, skipping full content retrieval")
            return relevant_items
            
        logger.info(f"Getting full content for {len(relevant_items)} relevant Wikipedia pages")
        
        results = []
        for item in relevant_items:
            title = item.get("id")  # Title stored as ID
            
            if not title:
                results.append(item)
                continue
                
            try:
                # Get the full page
                page = wikipedia.page(title, auto_suggest=False)
                
                # Create a full result with all information
                result = {
                    "title": page.title,
                    "link": page.url,
                    "snippet": item.get("snippet", "")  # Keep existing snippet
                }
                
                # Add additional information
                result["content"] = page.content
                result["full_content"] = page.content
                result["categories"] = page.categories
                result["references"] = page.references
                result["links"] = page.links
                result["images"] = page.images
                result["sections"] = page.sections
                
                results.append(result)
                
            except (wikipedia.exceptions.DisambiguationError, 
                   wikipedia.exceptions.PageError,
                   wikipedia.exceptions.WikipediaException) as e:
                # If error, use the preview
                logger.warning(f"Error getting full content for '{title}': {e}")
                results.append(item)
            except Exception as e:
                logger.error(f"Unexpected error getting full content for '{title}': {e}")
                results.append(item)
        
        return results
    
    def get_summary(self, title: str, sentences: Optional[int] = None) -> str:
        """
        Get a summary of a specific Wikipedia page.
        
        Args:
            title: Title of the Wikipedia page
            sentences: Number of sentences to include (defaults to self.sentences)
            
        Returns:
            Summary of the page
        """
        sentences = sentences or self.sentences
        try:
            return wikipedia.summary(title, sentences=sentences, auto_suggest=False)
        except wikipedia.exceptions.DisambiguationError as e:
            if e.options and len(e.options) > 0:
                return wikipedia.summary(e.options[0], sentences=sentences, auto_suggest=False)
            raise
    
    def get_page(self, title: str) -> Dict[str, Any]:
        """
        Get detailed information about a specific Wikipedia page.
        
        Args:
            title: Title of the Wikipedia page
            
        Returns:
            Dictionary with page information
        """
        # Check if we should include full content
        include_content = not (hasattr(config, 'SEARCH_SNIPPETS_ONLY') and config.SEARCH_SNIPPETS_ONLY)
        
        try:
            page = wikipedia.page(title, auto_suggest=False)
            
            result = {
                "title": page.title,
                "link": page.url,
                "snippet": self.get_summary(title, self.sentences)
            }
            
            # Add additional information if requested
            if include_content:
                result["content"] = page.content
                result["full_content"] = page.content
                result["categories"] = page.categories
                result["references"] = page.references
                result["links"] = page.links
                result["images"] = page.images
                result["sections"] = page.sections
            
            return result
        except wikipedia.exceptions.DisambiguationError as e:
            if e.options and len(e.options) > 0:
                return self.get_page(e.options[0])
            raise
    
    def set_language(self, language: str) -> None:
        """
        Change the Wikipedia language.
        
        Args:
            language: Language code (e.g., 'en', 'fr', 'es')
        """
        wikipedia.set_lang(language)