from typing import Any, Dict, List, Optional

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

from ...constants import SNIPPET_LENGTH_SHORT
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
    is_lexical = True
    needs_llm_relevance_filter = 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.
        # Runs as a preview filter (before LLM relevance) because Tiers 1-3
        # are instant data lookups — no point sending irrelevant journals
        # through the expensive LLM relevance filter.
        preview_filters = []
        journal_filter = self._create_journal_filter(
            "arxiv", llm, settings_snapshot
        )
        if journal_filter is not None:
            preview_filters.append(journal_filter)

        super().__init__(
            llm=llm,
            max_filtered_results=max_filtered_results,
            max_results=max_results,
            preview_filters=preview_filters,  # type: ignore[arg-type]
            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
        return list(client.results(search))

    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[:SNIPPET_LENGTH_SHORT] + "..."
                        if len(paper.summary) > SNIPPET_LENGTH_SHORT
                        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
        """
        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,
                        # Explicitly forward for journal quality filter
                        "journal_ref": paper.journal_ref,
                    }
                )

                # 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 pypdf first
                            try:
                                from pypdf import PdfReader

                                with open(paper_path, "rb") as pdf_file:
                                    pdf_reader = 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 pypdf"
                                        )
                            except (ImportError, Exception) as e1:
                                # Fall back to pdfplumber
                                try:
                                    import pdfplumber

                                    with pdfplumber.open(paper_path) as pdf:
                                        pdf_text = ""
                                        for plumber_page in pdf.pages:
                                            pdf_text += (
                                                plumber_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,
            }

            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
