from typing import Any, Dict, List, Optional

import wikipedia
from langchain_core.language_models import BaseLLM
from loguru import logger

from ...config import search_config
from ..search_engine_base import BaseSearchEngine


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:
            # Apply rate limiting before search request
            self._last_wait_time = self.rate_tracker.apply_rate_limit(
                self.engine_type
            )

            # 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:
                        # Apply rate limiting before summary request
                        self._last_wait_time = (
                            self.rate_tracker.apply_rate_limit(self.engine_type)
                        )

                        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.exception(
                                    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(' ', '_')}",
                            "source": "Wikipedia",
                        }

                        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:
                    logger.exception(f"Unexpected error for '{title}'")
                    continue

            logger.info(
                f"Successfully created {len(previews)} previews from Wikipedia"
            )
            return previews

        except Exception:
            logger.exception("Error getting Wikipedia previews")
            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(search_config, "SEARCH_SNIPPETS_ONLY")
            and search_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:
                # Apply rate limiting before page request
                self._last_wait_time = self.rate_tracker.apply_rate_limit(
                    self.engine_type
                )

                # 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
                    "source": "Wikipedia",
                }

                # 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.exception(
                    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
        """
        # Initialize include_content with our instance value
        include_content = self.include_content

        # Check if we should override with config setting
        if hasattr(search_config, "SEARCH_SNIPPETS_ONLY"):
            include_content = not search_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),
                "source": "Wikipedia",
            }

            # 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)
