"""
Paperless-ngx search engine implementation for Local Deep Research.

This module provides a proper search engine implementation that connects to a Paperless-ngx
instance, allowing LDR to search and retrieve documents from your personal
document management system.
"""

import re
from typing import Any, Dict, List, Optional
import requests
from urllib.parse import urljoin

from langchain_core.language_models import BaseLLM
from loguru import logger

from ..search_engine_base import BaseSearchEngine


class PaperlessSearchEngine(BaseSearchEngine):
    """Paperless-ngx search engine implementation with full LDR integration."""

    # Class constants for magic numbers
    MAX_SNIPPET_LENGTH = 3000  # Reasonable limit to avoid context window issues
    SNIPPET_CONTEXT_BEFORE = 500  # Characters before matched term in snippet
    SNIPPET_CONTEXT_AFTER = 2500  # Characters after matched term in snippet

    def __init__(
        self,
        api_url: str = None,
        api_key: str = None,
        api_token: str = None,  # Support both for backwards compatibility
        max_results: int = 10,
        timeout: int = 30,
        verify_ssl: bool = True,
        include_content: bool = True,
        llm: Optional[BaseLLM] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
        **kwargs,
    ):
        """
        Initialize the Paperless-ngx search engine.

        Args:
            api_url: Base URL of Paperless-ngx instance (e.g., "http://localhost:8000")
                    If not provided, will look for PAPERLESS_API_URL env var
            api_key: API token for authentication (preferred parameter name)
            api_token: API token for authentication (backwards compatibility)
                      If not provided, will look for PAPERLESS_API_TOKEN env var
            max_results: Maximum number of search results
            timeout: Request timeout in seconds
            verify_ssl: Whether to verify SSL certificates
            include_content: Whether to include document content in results
            llm: Language model for relevance filtering (optional)
            settings_snapshot: Settings snapshot for thread context
            **kwargs: Additional parameters passed to parent
        """
        super().__init__(
            max_results=max_results,
            llm=llm,
            settings_snapshot=settings_snapshot,
            **kwargs,
        )

        # Use provided configuration or get from settings
        self.api_url = api_url
        # Support both api_key and api_token for compatibility
        self.api_token = api_key or api_token

        # If no API URL provided, try to get from settings_snapshot
        if not self.api_url and settings_snapshot:
            self.api_url = settings_snapshot.get(
                "search.engine.web.paperless.default_params.api_url",
                "http://localhost:8000",
            )

        # If no API token provided, try to get from settings_snapshot
        if not self.api_token and settings_snapshot:
            self.api_token = settings_snapshot.get(
                "search.engine.web.paperless.api_key", ""
            )

        # Fix AttributeError: Check if api_url is None before calling rstrip
        if self.api_url:
            # Remove trailing slash from API URL
            self.api_url = self.api_url.rstrip("/")
        else:
            # Default to localhost if nothing provided
            self.api_url = "http://localhost:8000"
            logger.warning(
                "No Paperless API URL provided, using default: http://localhost:8000"
            )

        self.timeout = timeout
        self.verify_ssl = verify_ssl
        self.include_content = include_content

        # Set up headers for authentication
        self.headers = {}
        if self.api_token:
            self.headers["Authorization"] = f"Token {self.api_token}"

        logger.info(
            f"Initialized Paperless-ngx search engine for {self.api_url}"
        )

    def _make_request(
        self, endpoint: str, params: Optional[Dict] = None
    ) -> Dict[str, Any]:
        """
        Make a request to the Paperless-ngx API.

        Args:
            endpoint: API endpoint path
            params: Query parameters

        Returns:
            JSON response from the API
        """
        url = urljoin(self.api_url, endpoint)

        logger.debug(f"Making request to: {url}")
        logger.debug(f"Request params: {params}")
        logger.debug(
            f"Headers: {self.headers.keys() if self.headers else 'None'}"
        )

        try:
            response = requests.get(
                url,
                params=params,
                headers=self.headers,
                timeout=self.timeout,
                verify=self.verify_ssl,
            )
            response.raise_for_status()
            result = response.json()

            # Log response details
            if isinstance(result, dict):
                if "results" in result:
                    logger.info(
                        f"API returned {len(result.get('results', []))} results, total count: {result.get('count', 'unknown')}"
                    )
                    # Log first result details if available
                    if result.get("results"):
                        first = result["results"][0]
                        logger.debug(
                            f"First result: id={first.get('id')}, title='{first.get('title', 'No title')[:50]}...'"
                        )
                        if "__search_hit__" in first:
                            logger.debug(
                                f"Has search hit data with score={first['__search_hit__'].get('score')}"
                            )
                else:
                    logger.debug(f"API response keys: {result.keys()}")

            return result
        except requests.exceptions.RequestException:
            logger.exception("Error making request to Paperless-ngx")
            logger.debug(f"Failed URL: {url}, params: {params}")
            return {}

    def _expand_query_with_llm(self, query: str) -> str:
        """
        Use LLM to expand query with relevant keywords and synonyms.

        Args:
            query: Original search query

        Returns:
            Expanded query with keywords
        """
        if not self.llm:
            logger.info(
                f"No LLM available for query expansion, using original: '{query}'"
            )
            return query

        try:
            prompt = f"""Paperless-ngx uses TF-IDF keyword search, not semantic search.
Convert this query into keywords that would appear in documents.

Query: "{query}"

Output format: keyword1 OR keyword2 OR "multi word phrase" OR keyword3
Include synonyms, plural forms, and technical terms.

IMPORTANT: Output ONLY the search query. No explanations, no additional text."""

            logger.debug(
                f"Sending query expansion prompt to LLM for: '{query}'"
            )
            response = self.llm.invoke(prompt)
            expanded = response.content.strip()

            logger.debug(
                f"Raw LLM response (first 500 chars): {expanded[:500]}"
            )

            # Clean up the response - remove any explanatory text
            if "\n" in expanded:
                expanded = expanded.split("\n")[0]
                logger.debug("Took first line of LLM response")

            # Always trust the LLM's expansion - it knows better than hard-coded rules
            logger.info(
                f"LLM expanded query from '{query}' to {len(expanded)} chars with {expanded.count('OR')} ORs"
            )
            logger.debug(
                f"Expanded query preview (first 200 chars): {expanded[:200]}..."
            )
            return expanded

        except Exception:
            logger.exception("Failed to expand query with LLM")
            return query

    def _multi_pass_search(self, query: str) -> List[Dict[str, Any]]:
        """
        Perform multiple search passes with different strategies.

        Args:
            query: Original search query

        Returns:
            Combined and deduplicated results
        """
        logger.info(f"Starting multi-pass search for query: '{query}'")
        all_results = {}  # Use dict to deduplicate by doc_id

        # Pass 1: Original query
        params = {
            "query": query,
            "page_size": self.max_results,
            "ordering": "-score",
        }

        logger.info(
            f"Pass 1 - Original query: '{query}' (max_results={self.max_results})"
        )
        response = self._make_request("/api/documents/", params=params)

        if response and "results" in response:
            pass1_count = len(response["results"])
            logger.info(f"Pass 1 returned {pass1_count} documents")
            for doc in response["results"]:
                doc_id = doc.get("id")
                if doc_id and doc_id not in all_results:
                    all_results[doc_id] = doc
                    logger.debug(
                        f"Added doc {doc_id}: {doc.get('title', 'No title')}"
                    )
        else:
            logger.warning(
                f"Pass 1 returned no results or invalid response: {response}"
            )

        # Pass 2: LLM-expanded keywords (if LLM available)
        if self.llm:
            expanded_query = self._expand_query_with_llm(query)
            if expanded_query != query:
                params["query"] = expanded_query
                params["page_size"] = self.max_results * 2  # Get more results

                logger.info(
                    f"Pass 2 - Using expanded query with {expanded_query.count('OR')} ORs"
                )
                logger.debug(
                    f"Pass 2 - Full expanded query (first 500 chars): '{expanded_query[:500]}...'"
                )
                logger.info(
                    f"Pass 2 - Max results set to: {params['page_size']}"
                )
                response = self._make_request("/api/documents/", params=params)

                if response and "results" in response:
                    pass2_new = 0
                    for doc in response["results"]:
                        doc_id = doc.get("id")
                        if doc_id and doc_id not in all_results:
                            all_results[doc_id] = doc
                            pass2_new += 1
                            logger.debug(
                                f"Pass 2 added new doc {doc_id}: {doc.get('title', 'No title')}"
                            )
                    logger.info(
                        f"Pass 2 found {len(response['results'])} docs, added {pass2_new} new"
                    )
                else:
                    logger.warning("Pass 2 returned no results")
            else:
                logger.info("Pass 2 skipped - expanded query same as original")
        else:
            logger.info("Pass 2 skipped - no LLM available")

        # Sort by relevance score if available
        logger.info(f"Total unique documents collected: {len(all_results)}")
        sorted_results = sorted(
            all_results.values(),
            key=lambda x: x.get("__search_hit__", {}).get("score", 0),
            reverse=True,
        )

        final_results = sorted_results[: self.max_results]
        logger.info(
            f"Returning top {len(final_results)} documents after sorting by score"
        )

        # Log titles and scores of final results
        for i, doc in enumerate(final_results[:5], 1):  # Log first 5
            score = doc.get("__search_hit__", {}).get("score", 0)
            logger.debug(
                f"Result {i}: '{doc.get('title', 'No title')}' (score={score})"
            )

        return final_results

    def _get_previews(self, query: str) -> List[Dict[str, Any]]:
        """
        Get preview results from Paperless-ngx using multi-pass strategy.

        Args:
            query: Search query

        Returns:
            List of preview dictionaries
        """
        try:
            # Use multi-pass search strategy
            results = self._multi_pass_search(query)

            if not results:
                return []

            # Convert documents to preview format
            # Note: Each document may return multiple previews (one per highlight)
            previews = []
            for doc_data in results:
                doc_previews = self._convert_document_to_preview(
                    doc_data, query
                )
                # Handle both single preview and list of previews
                if isinstance(doc_previews, list):
                    previews.extend(doc_previews)
                else:
                    previews.append(doc_previews)

            logger.info(
                f"Found {len(previews)} documents in Paperless-ngx for query: {query}"
            )
            return previews

        except Exception:
            logger.exception("Error getting previews from Paperless-ngx")
            return []

    def _convert_document_to_preview(
        self, doc_data: Dict[str, Any], query: str = ""
    ) -> Dict[str, Any]:
        """
        Convert a Paperless-ngx document to LDR preview format.

        Args:
            doc_data: Document data from the API
            query: Original search query (for context)

        Returns:
            Preview dictionary in LDR format
        """
        # Extract title
        title = doc_data.get("title", f"Document {doc_data.get('id')}")
        doc_id = doc_data.get("id")

        logger.info(
            f"Converting document {doc_id}: '{title}' to preview format"
        )

        # Build URL - use the web interface URL for user access
        url = f"{self.api_url}/documents/{doc_id}/details"
        logger.debug(f"Generated URL for doc {doc_id}: {url}")

        # Extract snippet - prefer highlighted content from search
        snippet = ""
        search_score = 0.0
        search_rank = None
        all_highlights = []  # Initialize empty highlights list

        if "__search_hit__" in doc_data:
            search_hit = doc_data["__search_hit__"]
            logger.debug(
                f"Found __search_hit__ data for doc {doc_id}: score={search_hit.get('score')}, rank={search_hit.get('rank')}"
            )

            # Get highlights - this is the search snippet with matched terms
            if search_hit.get("highlights"):
                # Highlights can be a string or list
                highlights = search_hit.get("highlights")
                logger.info(
                    f"Found highlights for doc {doc_id}: type={type(highlights).__name__}, length={len(str(highlights))}"
                )

                if isinstance(highlights, list):
                    logger.debug(
                        f"Highlights is list with {len(highlights)} items"
                    )
                    # IMPORTANT: Store highlights list for processing later
                    # Each highlight will become a separate search result for proper citation
                    all_highlights = highlights
                    # Use first highlight for the default snippet
                    snippet = highlights[0] if highlights else ""
                    logger.info(
                        f"Will create {len(highlights)} separate results from highlights"
                    )
                else:
                    all_highlights = [
                        str(highlights)
                    ]  # Single highlight as list
                    snippet = str(highlights)

                logger.debug(
                    f"Raw snippet before cleaning (first 200 chars): {snippet[:200]}"
                )

                # Clean HTML tags but preserve the matched text
                snippet = re.sub(r"<span[^>]*>", "**", snippet)
                snippet = re.sub(r"</span>", "**", snippet)
                snippet = re.sub(r"<[^>]+>", "", snippet)

                logger.debug(
                    f"Cleaned snippet (first 200 chars): {snippet[:200]}"
                )

                # Limit snippet length to avoid context window issues
                if (
                    self.MAX_SNIPPET_LENGTH
                    and len(snippet) > self.MAX_SNIPPET_LENGTH
                ):
                    # Cut at word boundary to avoid mid-word truncation
                    snippet = (
                        snippet[: self.MAX_SNIPPET_LENGTH].rsplit(" ", 1)[0]
                        + "..."
                    )
                    logger.debug(
                        f"Truncated snippet to {self.MAX_SNIPPET_LENGTH} chars"
                    )

            # Get search relevance metadata
            search_score = search_hit.get("score", 0.0)
            search_rank = search_hit.get("rank")
            logger.info(
                f"Search metadata for doc {doc_id}: score={search_score}, rank={search_rank}"
            )
        else:
            logger.warning(
                f"No __search_hit__ data for doc {doc_id}, will use content fallback"
            )

        if not snippet:
            logger.info(
                f"No snippet from highlights for doc {doc_id}, using content fallback"
            )
            # Fallback to content preview if no highlights available
            content = doc_data.get("content", "")
            if content:
                logger.debug(f"Document has content of length {len(content)}")
                # Try to find context around query terms if possible
                if query:
                    query_terms = query.lower().split()
                    content_lower = content.lower()
                    logger.debug(
                        f"Searching for query terms in content: {query_terms}"
                    )

                    # Find first occurrence of any query term
                    best_pos = -1
                    for term in query_terms:
                        pos = content_lower.find(term)
                        if pos != -1 and (best_pos == -1 or pos < best_pos):
                            best_pos = pos
                            logger.debug(
                                f"Found term '{term}' at position {pos}"
                            )

                    if best_pos != -1:
                        # Extract context around the found term - much larger context for research
                        start = max(0, best_pos - 2000)
                        end = min(len(content), best_pos + 8000)
                        snippet = "..." + content[start:end] + "..."
                        logger.info(
                            f"Extracted snippet around query term at position {best_pos}"
                        )
                    else:
                        # Just take the beginning - use 10000 chars for research
                        snippet = content[:10000]
                        logger.info(
                            "No query terms found, using first 10000 chars of content"
                        )
                else:
                    snippet = content[:10000]
                    logger.info(
                        "No query provided, using first 10000 chars of content"
                    )

                if len(content) > 10000:
                    snippet += "..."
            else:
                logger.warning(f"No content available for doc {doc_id}")

        logger.info(f"Final snippet for doc {doc_id} has length {len(snippet)}")

        # Build metadata
        metadata = {
            "doc_id": str(doc_id),
            "correspondent": doc_data.get("correspondent_name", ""),
            "document_type": doc_data.get("document_type_name", ""),
            "created": doc_data.get("created", ""),
            "modified": doc_data.get("modified", ""),
            "archive_serial_number": doc_data.get("archive_serial_number"),
            "search_score": search_score,
            "search_rank": search_rank,
        }

        # Add tags if present
        tags = doc_data.get("tags_list", [])
        if isinstance(tags, list) and tags:
            metadata["tags"] = ", ".join(str(tag) for tag in tags)

        # Build enhanced title with available metadata for better citations
        title_parts = []

        # Add correspondent/author if available
        correspondent = doc_data.get("correspondent_name", "")
        if correspondent:
            title_parts.append(f"{correspondent}.")
            logger.debug(f"Added correspondent to title: {correspondent}")

        # Add the document title
        title_parts.append(title)

        # Add document type if it's meaningful (not just generic types)
        doc_type = doc_data.get("document_type_name", "")
        if doc_type and doc_type not in ["Letter", "Other", "Document", ""]:
            title_parts.append(f"({doc_type})")
            logger.debug(f"Added document type to title: {doc_type}")

        # Add year from created date if available
        created_date = doc_data.get("created", "")
        if created_date and len(created_date) >= 4:
            year = created_date[:4]
            title_parts.append(year)
            logger.debug(f"Added year to title: {year}")

        # Format the enhanced title for display in sources list
        if title_parts:
            enhanced_title = " ".join(title_parts)
        else:
            enhanced_title = title

        logger.info(f"Enhanced title for doc {doc_id}: '{enhanced_title}'")

        # Build the preview
        preview = {
            "title": enhanced_title,  # Use enhanced title with bibliographic info
            "url": url,
            "link": url,  # Add 'link' key for compatibility with search utilities
            "snippet": snippet,
            "author": doc_data.get("correspondent_name", ""),
            "date": doc_data.get("created", ""),
            "source": "Paperless",  # Keep source as the system name like other engines
            "metadata": metadata,
            "_raw_data": doc_data,  # Store raw data for full content retrieval
        }

        logger.info(
            f"Built preview for doc {doc_id}: URL={url}, snippet_len={len(snippet)}, has_author={bool(preview['author'])}, has_date={bool(preview['date'])}"
        )

        # Check if we have multiple highlights to return as separate results
        if len(all_highlights) > 1:
            # Create multiple previews, one for each highlight
            previews = []
            for i, highlight in enumerate(all_highlights):
                # Clean each highlight
                clean_snippet = re.sub(r"<span[^>]*>", "**", str(highlight))
                clean_snippet = re.sub(r"</span>", "**", clean_snippet)
                clean_snippet = re.sub(r"<[^>]+>", "", clean_snippet)

                # Create a preview for this highlight
                highlight_preview = {
                    "title": f"{enhanced_title} (excerpt {i + 1})",  # Differentiate each excerpt
                    "url": url,
                    "link": url,
                    "snippet": clean_snippet,
                    "author": doc_data.get("correspondent_name", ""),
                    "date": doc_data.get("created", ""),
                    "source": "Paperless",
                    "metadata": {
                        **metadata,
                        "excerpt_number": i + 1,
                        "total_excerpts": len(all_highlights),
                    },
                    "_raw_data": doc_data,
                }
                previews.append(highlight_preview)

            logger.info(
                f"Created {len(previews)} separate previews from highlights for doc {doc_id}"
            )
            return previews
        else:
            # Single preview (original behavior)
            return preview

    def _get_full_content(
        self, relevant_items: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """
        Get full content for relevant documents.

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of dictionaries with full content
        """
        if not self.include_content:
            # If content inclusion is disabled, just return previews
            return relevant_items

        logger.info(f"Getting full content for {len(relevant_items)} documents")
        results = []
        for idx, item in enumerate(relevant_items):
            try:
                logger.info(
                    f"Processing document {idx + 1}: title='{item.get('title', 'No title')[:50]}...', url={item.get('url', 'No URL')}"
                )
                logger.debug(f"Document {idx + 1} keys: {item.keys()}")
                logger.debug(
                    f"Document {idx + 1} has snippet of length: {len(item.get('snippet', ''))}"
                )

                # Get the full document content if we have the raw data
                if "_raw_data" in item:
                    doc_data = item["_raw_data"]
                    full_content = doc_data.get("content", "")

                    if not full_content:
                        # Try to fetch the document details
                        doc_id = item["metadata"].get("doc_id")
                        if doc_id:
                            detail_response = self._make_request(
                                f"/api/documents/{doc_id}/"
                            )
                            if detail_response:
                                full_content = detail_response.get(
                                    "content", ""
                                )

                    item["full_content"] = full_content or item["snippet"]
                    logger.info(
                        f"Document {idx + 1} full content length: {len(item['full_content'])}"
                    )
                else:
                    # Fallback to snippet if no raw data
                    item["full_content"] = item["snippet"]
                    logger.info(
                        f"Document {idx + 1} using snippet as full content (no raw data)"
                    )

                # Log the final document structure for debugging citation issues
                logger.info(
                    f"Document {idx + 1} final structure: title='{item.get('title', '')[:50]}...', has_link={bool(item.get('link'))}, has_url={bool(item.get('url'))}, source='{item.get('source', 'Unknown')}'"
                )

                # Remove the raw data from the result
                item.pop("_raw_data", None)
                results.append(item)

            except Exception:
                logger.exception("Error getting full content for document")
                item["full_content"] = item["snippet"]
                item.pop("_raw_data", None)
                results.append(item)

        return results

    def run(
        self, query: str, research_context: Dict[str, Any] | None = None
    ) -> List[Dict[str, Any]]:
        """
        Execute search on Paperless-ngx.

        Args:
            query: Search query
            research_context: Context from previous research

        Returns:
            List of search results in LDR format
        """
        try:
            # Get previews
            previews = self._get_previews(query)

            if not previews:
                return []

            # Apply LLM filtering if available
            if (
                self.llm
                and hasattr(self, "_content_filters")
                and self._content_filters
            ):
                filtered_previews = self._apply_content_filters(previews, query)
            else:
                filtered_previews = previews

            # Get full content for relevant items
            results = self._get_full_content(filtered_previews)

            logger.info(
                f"Search completed successfully, returning {len(results)} results"
            )
            # Enhanced logging to track document structure for citation debugging
            for i, r in enumerate(results[:3], 1):
                logger.info(
                    f"Result {i}: title='{r.get('title', '')[:50]}...', "
                    f"has_full_content={bool(r.get('full_content'))}, "
                    f"full_content_len={len(r.get('full_content', ''))}, "
                    f"snippet_len={len(r.get('snippet', ''))}, "
                    f"url={r.get('url', '')[:50]}"
                )

            return results

        except Exception:
            logger.exception("Error in Paperless-ngx search")
            return []

    async def arun(self, query: str) -> List[Dict[str, Any]]:
        """
        Async version of search.

        Currently falls back to sync version.
        """
        return self.run(query)

    def test_connection(self) -> bool:
        """
        Test the connection to Paperless-ngx.

        Returns:
            True if connection successful, False otherwise
        """
        try:
            response = self._make_request("/api/")
            return bool(response)
        except Exception:
            logger.exception("Failed to connect to Paperless-ngx")
            return False

    def get_document_count(self) -> int:
        """
        Get the total number of documents in Paperless-ngx.

        Returns:
            Number of documents, or -1 if error
        """
        try:
            response = self._make_request(
                "/api/documents/", params={"page_size": 1}
            )
            return response.get("count", -1)
        except Exception:
            return -1
