from typing import Any, Dict, List, Optional

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

from ...advanced_search_system.filters.journal_reputation_filter import (
    JournalReputationFilter,
)
from ...config import search_config
from ..rate_limiting import RateLimitError
from ..search_engine_base import BaseSearchEngine


class ArXivSearchEngine(BaseSearchEngine):
    """arXiv search engine implementation with two-phase approach"""

    # Mark as public search engine
    is_public = True
    # Not a generic search engine (specialized for academic papers)
    is_generic = False
    # Scientific/academic search engine
    is_scientific = True

    def __init__(
        self,
        max_results: int = 10,
        sort_by: str = "relevance",
        sort_order: str = "descending",
        include_full_text: bool = False,
        download_dir: Optional[str] = None,
        max_full_text: int = 1,
        llm: Optional[BaseLLM] = None,
        max_filtered_results: Optional[int] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
    ):  # Added this parameter
        """
        Initialize the arXiv search engine.

        Args:
            max_results: Maximum number of search results
            sort_by: Sorting criteria ('relevance', 'lastUpdatedDate', or 'submittedDate')
            sort_order: Sort order ('ascending' or 'descending')
            include_full_text: Whether to include full paper content in results (downloads PDF)
            download_dir: Directory to download PDFs to (if include_full_text is True)
            max_full_text: Maximum number of PDFs to download and process (default: 1)
            llm: Language model for relevance filtering
            max_filtered_results: Maximum number of results to keep after filtering
            settings_snapshot: Settings snapshot for thread context
        """
        # Initialize the journal reputation filter if needed.
        content_filters = []
        journal_filter = JournalReputationFilter.create_default(
            model=llm, engine_name="arxiv", settings_snapshot=settings_snapshot
        )
        if journal_filter is not None:
            content_filters.append(journal_filter)

        # 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,
            # We deliberately do this filtering after relevancy checks,
            # because it is potentially quite slow.
            content_filters=content_filters,
            settings_snapshot=settings_snapshot,
        )
        self.max_results = max(self.max_results, 25)
        self.sort_by = sort_by
        self.sort_order = sort_order
        self.include_full_text = include_full_text
        self.download_dir = download_dir
        self.max_full_text = max_full_text

        # Map sort parameters to arxiv package parameters
        self.sort_criteria = {
            "relevance": arxiv.SortCriterion.Relevance,
            "lastUpdatedDate": arxiv.SortCriterion.LastUpdatedDate,
            "submittedDate": arxiv.SortCriterion.SubmittedDate,
        }

        self.sort_directions = {
            "ascending": arxiv.SortOrder.Ascending,
            "descending": arxiv.SortOrder.Descending,
        }

    def _get_search_results(self, query: str) -> List[Any]:
        """
        Helper method to get search results from arXiv API.

        Args:
            query: The search query

        Returns:
            List of arXiv paper objects
        """
        # Configure the search client
        sort_criteria = self.sort_criteria.get(
            self.sort_by, arxiv.SortCriterion.Relevance
        )
        sort_order = self.sort_directions.get(
            self.sort_order, arxiv.SortOrder.Descending
        )

        # Create the search client
        client = arxiv.Client(page_size=self.max_results)

        # Create the search query
        search = arxiv.Search(
            query=query,
            max_results=self.max_results,
            sort_by=sort_criteria,
            sort_order=sort_order,
        )

        # Apply rate limiting before making the request
        self._last_wait_time = self.rate_tracker.apply_rate_limit(
            self.engine_type
        )

        # Get the search results
        papers = list(client.results(search))

        return papers

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview information for arXiv papers.

        Args:
            query: The search query

        Returns:
            List of preview dictionaries
        """
        logger.info("Getting paper previews from arXiv")

        try:
            # Get search results from arXiv
            papers = self._get_search_results(query)

            # Store the paper objects for later use
            self._papers = {paper.entry_id: paper for paper in papers}

            # Format results as previews with basic information
            previews = []
            for paper in papers:
                preview = {
                    "id": paper.entry_id,  # Use entry_id as ID
                    "title": paper.title,
                    "link": paper.entry_id,  # arXiv URL
                    "snippet": (
                        paper.summary[:250] + "..."
                        if len(paper.summary) > 250
                        else paper.summary
                    ),
                    "authors": [
                        author.name for author in paper.authors[:3]
                    ],  # First 3 authors
                    "published": (
                        paper.published.strftime("%Y-%m-%d")
                        if paper.published
                        else None
                    ),
                    "journal_ref": paper.journal_ref,
                    "source": "arXiv",
                }

                previews.append(preview)

            return previews

        except Exception as e:
            error_msg = str(e)
            logger.exception("Error getting arXiv previews")

            # Check for rate limiting patterns
            if (
                "429" in error_msg
                or "too many requests" in error_msg.lower()
                or "rate limit" in error_msg.lower()
                or "service unavailable" in error_msg.lower()
                or "503" in error_msg
            ):
                raise RateLimitError(f"arXiv rate limit hit: {error_msg}")

            return []

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for the relevant arXiv papers.
        Downloads PDFs and extracts text when include_full_text is True.
        Limits the number of PDFs processed to max_full_text.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content
        """
        # Check if we should get 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("Getting full content for relevant arXiv papers")

        results = []
        pdf_count = 0  # Track number of PDFs processed

        for item in relevant_items:
            # Start with the preview data
            result = item.copy()

            # Get the paper ID
            paper_id = item.get("id")

            # Try to get the full paper from our cache
            paper = None
            if hasattr(self, "_papers") and paper_id in self._papers:
                paper = self._papers[paper_id]

            if paper:
                # Add complete paper information
                result.update(
                    {
                        "pdf_url": paper.pdf_url,
                        "authors": [
                            author.name for author in paper.authors
                        ],  # All authors
                        "published": (
                            paper.published.strftime("%Y-%m-%d")
                            if paper.published
                            else None
                        ),
                        "updated": (
                            paper.updated.strftime("%Y-%m-%d")
                            if paper.updated
                            else None
                        ),
                        "categories": paper.categories,
                        "summary": paper.summary,  # Full summary
                        "comment": paper.comment,
                        "doi": paper.doi,
                    }
                )

                # Default to using summary as content
                result["content"] = paper.summary
                result["full_content"] = paper.summary

                # Download PDF and extract text if requested and within limit
                if (
                    self.include_full_text
                    and self.download_dir
                    and pdf_count < self.max_full_text
                ):
                    try:
                        # Download the paper
                        pdf_count += (
                            1  # Increment counter before attempting download
                        )
                        # Apply rate limiting before PDF download
                        self.rate_tracker.apply_rate_limit(self.engine_type)

                        paper_path = paper.download_pdf(
                            dirpath=self.download_dir
                        )
                        result["pdf_path"] = str(paper_path)

                        # Extract text from PDF
                        try:
                            # Try PyPDF2 first
                            try:
                                import PyPDF2

                                with open(paper_path, "rb") as pdf_file:
                                    pdf_reader = PyPDF2.PdfReader(pdf_file)
                                    pdf_text = ""
                                    for page in pdf_reader.pages:
                                        pdf_text += page.extract_text() + "\n\n"

                                    if (
                                        pdf_text.strip()
                                    ):  # Only use if we got meaningful text
                                        result["content"] = pdf_text
                                        result["full_content"] = pdf_text
                                        logger.info(
                                            "Successfully extracted text from PDF using PyPDF2"
                                        )
                            except (ImportError, Exception) as e1:
                                # Fall back to pdfplumber
                                try:
                                    import pdfplumber

                                    with pdfplumber.open(paper_path) as pdf:
                                        pdf_text = ""
                                        for page in pdf.pages:
                                            pdf_text += (
                                                page.extract_text() + "\n\n"
                                            )

                                        if (
                                            pdf_text.strip()
                                        ):  # Only use if we got meaningful text
                                            result["content"] = pdf_text
                                            result["full_content"] = pdf_text
                                            logger.info(
                                                "Successfully extracted text from PDF using pdfplumber"
                                            )
                                except (ImportError, Exception) as e2:
                                    logger.exception(
                                        f"PDF text extraction failed: {e1!s}, then {e2!s}"
                                    )
                                    logger.info(
                                        "Using paper summary as content instead"
                                    )
                        except Exception:
                            logger.exception("Error extracting text from PDF")
                            logger.info(
                                "Using paper summary as content instead"
                            )
                    except Exception:
                        logger.exception(
                            f"Error downloading paper {paper.title}"
                        )
                        result["pdf_path"] = None
                        pdf_count -= 1  # Decrement counter if download fails
                elif (
                    self.include_full_text
                    and self.download_dir
                    and pdf_count >= self.max_full_text
                ):
                    # Reached PDF limit
                    logger.info(
                        f"Maximum number of PDFs ({self.max_full_text}) reached. Skipping remaining PDFs."
                    )
                    result["content"] = paper.summary
                    result["full_content"] = paper.summary

            results.append(result)

        return results

    def run(
        self, query: str, research_context: Dict[str, Any] | None = None
    ) -> List[Dict[str, Any]]:
        """
        Execute a search using arXiv with the two-phase approach.

        Args:
            query: The search query
            research_context: Context from previous research to use.

        Returns:
            List of search results
        """
        logger.info("---Execute a search using arXiv---")

        # Use the implementation from the parent class which handles all phases
        results = super().run(query, research_context=research_context)

        # Clean up
        if hasattr(self, "_papers"):
            del self._papers

        return results

    def get_paper_details(self, arxiv_id: str) -> Dict[str, Any]:
        """
        Get detailed information about a specific arXiv paper.

        Args:
            arxiv_id: arXiv ID of the paper (e.g., '2101.12345')

        Returns:
            Dictionary with paper information
        """
        try:
            # Create the search client
            client = arxiv.Client()

            # Search for the specific paper
            search = arxiv.Search(id_list=[arxiv_id], max_results=1)

            # Apply rate limiting before fetching paper by ID
            self._last_wait_time = self.rate_tracker.apply_rate_limit(
                self.engine_type
            )

            # Get the paper
            papers = list(client.results(search))
            if not papers:
                return {}

            paper = papers[0]

            # Format result based on config
            result = {
                "title": paper.title,
                "link": paper.entry_id,
                "snippet": (
                    paper.summary[:250] + "..."
                    if len(paper.summary) > 250
                    else paper.summary
                ),
                "authors": [
                    author.name for author in paper.authors[:3]
                ],  # First 3 authors
                "journal_ref": paper.journal_ref,
            }

            # Add full content if not in snippet-only mode
            if (
                not hasattr(search_config, "SEARCH_SNIPPETS_ONLY")
                or not search_config.SEARCH_SNIPPETS_ONLY
            ):
                result.update(
                    {
                        "pdf_url": paper.pdf_url,
                        "authors": [
                            author.name for author in paper.authors
                        ],  # All authors
                        "published": (
                            paper.published.strftime("%Y-%m-%d")
                            if paper.published
                            else None
                        ),
                        "updated": (
                            paper.updated.strftime("%Y-%m-%d")
                            if paper.updated
                            else None
                        ),
                        "categories": paper.categories,
                        "summary": paper.summary,  # Full summary
                        "comment": paper.comment,
                        "doi": paper.doi,
                        "content": paper.summary,  # Use summary as content
                        "full_content": paper.summary,  # For consistency
                    }
                )

                # Download PDF if requested
                if self.include_full_text and self.download_dir:
                    try:
                        # Apply rate limiting before PDF download
                        self.rate_tracker.apply_rate_limit(self.engine_type)

                        # Download the paper
                        paper_path = paper.download_pdf(
                            dirpath=self.download_dir
                        )
                        result["pdf_path"] = str(paper_path)
                    except Exception:
                        logger.exception("Error downloading paper")

            return result

        except Exception:
            logger.exception("Error getting paper details")
            return {}

    def search_by_author(
        self, author_name: str, max_results: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """
        Search for papers by a specific author.

        Args:
            author_name: Name of the author
            max_results: Maximum number of results (defaults to self.max_results)

        Returns:
            List of papers by the author
        """
        original_max_results = self.max_results

        try:
            if max_results:
                self.max_results = max_results

            query = f'au:"{author_name}"'
            return self.run(query)

        finally:
            # Restore original value
            self.max_results = original_max_results

    def search_by_category(
        self, category: str, max_results: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """
        Search for papers in a specific arXiv category.

        Args:
            category: arXiv category (e.g., 'cs.AI', 'physics.optics')
            max_results: Maximum number of results (defaults to self.max_results)

        Returns:
            List of papers in the category
        """
        original_max_results = self.max_results

        try:
            if max_results:
                self.max_results = max_results

            query = f"cat:{category}"
            return self.run(query)

        finally:
            # Restore original value
            self.max_results = original_max_results
