from loguru import logger
from typing import Any, Dict, List, Optional

from elasticsearch import Elasticsearch
from langchain_core.language_models import BaseLLM

from ...constants import SNIPPET_LENGTH_SHORT
from ..search_engine_base import BaseSearchEngine
from ...constants import DEFAULT_SEARCH_TOOL


class ElasticsearchSearchEngine(BaseSearchEngine):
    """Elasticsearch search engine implementation with two-phase approach"""

    is_local = True
    is_lexical = True
    needs_llm_relevance_filter = True
    # secrets to redact from error messages (see BaseSearchEngine._scrub_error)
    _secret_attrs = ("_api_key", "_password")
    # url_setting feeds the PDP's fail-up URL override: when the configured
    # hosts resolve to a PUBLIC endpoint (e.g. Elastic Cloud), the engine is
    # reclassified public so PRIVATE_ONLY denies it at selection time —
    # queries would leave the box even though the DATA is "local" in nature.
    # A localhost ES keeps the static is_local classification above.
    # cloud_id (which is NOT a host the PDP can classify) is handled
    # separately in __init__: it is rejected when the effective scope forbids
    # public egress.
    url_setting = "search.engine.web.elasticsearch.default_params.hosts"

    @staticmethod
    def _cloud_id_forbidden_by_scope(
        settings_snapshot: Optional[Dict[str, Any]],
    ) -> bool:
        """True when the effective egress scope forbids the public Elastic
        Cloud endpoint a ``cloud_id`` targets.

        Resolves the scope (including ADAPTIVE) via ``context_from_snapshot``
        and returns True for PRIVATE_ONLY / STRICT. Fails CLOSED (forbidden)
        if the policy cannot be evaluated, so a snapshot/policy error cannot
        open a cloud egress under a private posture. A missing/empty snapshot
        resolves to the permissive default (BOTH) and is allowed.
        """
        try:
            from ...security.egress.policy import (
                EgressScope,
                context_from_snapshot,
            )
            from ...config.thread_settings import get_setting_from_snapshot

            snapshot = settings_snapshot or {}
            primary = (
                get_setting_from_snapshot(
                    "search.tool",
                    default=DEFAULT_SEARCH_TOOL,
                    settings_snapshot=snapshot,
                )
                or DEFAULT_SEARCH_TOOL
            )
            ctx = context_from_snapshot(snapshot, primary)
            return ctx.scope in (
                EgressScope.PRIVATE_ONLY,
                EgressScope.STRICT,
            )
        except Exception:
            logger.bind(policy_audit=True).warning(
                "elasticsearch cloud_id egress check failed; failing closed",
                exc_info=True,
            )
            return True

    def __init__(
        self,
        hosts: Optional[List[str]] = None,
        index_name: str = "documents",
        username: Optional[str] = None,
        password: Optional[str] = None,
        api_key: Optional[str] = None,
        cloud_id: Optional[str] = None,
        max_results: int = 10,
        highlight_fields: List[str] = ["content", "title"],
        search_fields: List[str] = ["content", "title"],
        filter_query: Optional[Dict[str, Any]] = None,
        llm: Optional[BaseLLM] = None,
        max_filtered_results: Optional[int] = None,
        settings_snapshot: Optional[Dict[str, Any]] = None,
    ):
        """
        Initialize the Elasticsearch search engine.

        Args:
            hosts: List of Elasticsearch hosts
            index_name: Name of the index to search
            username: Optional username for authentication
            password: Optional password for authentication
            api_key: Optional API key for authentication
            cloud_id: Optional Elastic Cloud ID
            max_results: Maximum number of search results
            highlight_fields: Fields to highlight in search results
            search_fields: Fields to search in
            filter_query: Optional filter query in Elasticsearch DSL format
            llm: Language model for relevance filtering
            max_filtered_results: Maximum number of results to keep after filtering
        """
        # 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,
            settings_snapshot=settings_snapshot,
        )

        self.index_name = index_name
        self.highlight_fields = self._ensure_list(
            highlight_fields, default=["content", "title"]
        )
        self.search_fields = self._ensure_list(
            search_fields, default=["content", "title"]
        )
        self.filter_query = filter_query or {}

        # Store credentials for error-message redaction
        self._api_key = api_key
        self._password = password

        # Normalize hosts – may arrive as a JSON-encoded string from settings
        hosts = self._ensure_list(hosts, default=["http://localhost:9200"])

        # Initialize the Elasticsearch client
        es_args: Dict[str, Any] = {}

        # Basic authentication
        if username and password:
            es_args["basic_auth"] = (username, password)

        # API key authentication
        if api_key:
            es_args["api_key"] = api_key

        # Cloud ID for Elastic Cloud
        if cloud_id:
            # Egress policy: a cloud_id always targets a public Elastic Cloud
            # endpoint (*.cloud.es.io), but the url_setting reclassification
            # only inspects `hosts`. A cloud_id-only config would otherwise
            # keep the engine's static is_local=True and slip past
            # evaluate_engine, then connect at self.client.info() below. Reject
            # it when the effective scope forbids public egress (fail closed).
            if self._cloud_id_forbidden_by_scope(settings_snapshot):
                from ...security.egress.policy import (
                    Decision,
                    PolicyDeniedError,
                )

                logger.bind(policy_audit=True).warning(
                    "refusing Elasticsearch cloud_id under private egress scope"
                )
                raise PolicyDeniedError(
                    Decision(False, "elasticsearch_cloud_id_public_egress"),
                    target="search_engine:elasticsearch",
                )
            es_args["cloud_id"] = cloud_id

        # Connect to Elasticsearch
        self.client = Elasticsearch(hosts, **es_args)

        # Verify connection
        try:
            info = self.client.info()
            logger.info(
                f"Connected to Elasticsearch cluster: {info.get('cluster_name')}"
            )
            logger.info(
                f"Elasticsearch version: {info.get('version', {}).get('number')}"
            )
        except Exception as e:
            safe_msg = self._scrub_error(e)
            logger.warning(f"Failed to connect to Elasticsearch: {safe_msg}")
            raise ConnectionError(
                f"Could not connect to Elasticsearch: {safe_msg}"
            ) from None

    def close(self) -> None:
        """Close the Elasticsearch client and its connection pool."""
        from ...utilities.resource_utils import safe_close

        safe_close(self.client, "Elasticsearch client")
        super().close()

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

        Args:
            query: The search query

        Returns:
            List of preview dictionaries
        """
        logger.info(
            f"Getting document previews from Elasticsearch with query: {query}"
        )

        try:
            # Build the search query
            search_query = {
                "query": {
                    "multi_match": {
                        "query": query,
                        "fields": self.search_fields,
                        "type": "best_fields",
                        "tie_breaker": 0.3,
                    }
                },
                "highlight": {
                    "fields": {field: {} for field in self.highlight_fields},
                    "pre_tags": ["<em>"],
                    "post_tags": ["</em>"],
                },
                "size": self.max_results,
            }

            # Add filter if provided
            if self.filter_query:
                search_query["query"] = {
                    "bool": {
                        "must": search_query["query"],
                        "filter": self.filter_query,
                    }
                }

            # Execute the search
            response = self.client.search(
                index=self.index_name,
                body=search_query,
            )

            # Process the search results
            hits = response.get("hits", {}).get("hits", [])

            # Format results as previews with basic information
            previews = []
            for hit in hits:
                source = hit.get("_source", {})
                highlight = hit.get("highlight", {})

                # Extract highlighted snippets or fall back to original content
                snippet = ""
                for field in self.highlight_fields:
                    if highlight.get(field):
                        # Join all highlights for this field
                        field_snippets = " ... ".join(highlight[field])
                        snippet += field_snippets + " "

                # If no highlights, use a portion of the content
                if not snippet and "content" in source:
                    content = source.get("content", "")
                    snippet = (
                        content[:SNIPPET_LENGTH_SHORT] + "..."
                        if len(content) > SNIPPET_LENGTH_SHORT
                        else content
                    )

                # Create preview object
                preview = {
                    "id": hit.get("_id", ""),
                    "title": source.get("title", "Untitled Document"),
                    "link": source.get("url", "")
                    or f"elasticsearch://{self.index_name}/{hit.get('_id', '')}",
                    "snippet": snippet.strip(),
                    "score": hit.get("_score", 0),
                    "_index": hit.get("_index", self.index_name),
                }

                previews.append(preview)

            logger.info(
                f"Found {len(previews)} preview results from Elasticsearch"
            )
            return previews

        except Exception as e:
            safe_msg = self._scrub_error(e)
            logger.warning(f"Error getting Elasticsearch previews: {safe_msg}")
            return []

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

        Args:
            relevant_items: List of relevant preview dictionaries

        Returns:
            List of result dictionaries with full content
        """
        logger.info("Getting full content for relevant Elasticsearch documents")

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

            # Get the document ID
            doc_id = item.get("id")
            if not doc_id:
                # Skip items without ID
                logger.warning(f"Skipping item without ID: {item}")
                results.append(result)
                continue

            try:
                # Fetch the full document
                doc_response = self.client.get(
                    index=self.index_name,
                    id=doc_id,
                )

                # Get the source document
                source = doc_response.get("_source", {})

                # Add full content to the result
                result["content"] = source.get(
                    "content", result.get("snippet", "")
                )
                result["full_content"] = source.get("content", "")

                # Add metadata from source
                for key, value in source.items():
                    if key not in result and key not in ["content"]:
                        result[key] = value

            except Exception as e:
                safe_msg = self._scrub_error(e)
                logger.warning(
                    f"Error fetching full content for document {doc_id}: {safe_msg}"
                )
                # Keep the preview data if we can't get the full content

            results.append(result)

        return results

    def search_by_query_string(self, query_string: str) -> List[Dict[str, Any]]:
        """
        Perform a search using Elasticsearch Query String syntax.

        Args:
            query_string: The query in Elasticsearch Query String syntax

        Returns:
            List of search results
        """
        try:
            # Build the search query
            search_query = {
                "query": {
                    "query_string": {
                        "query": query_string,
                        "fields": self.search_fields,
                    }
                },
                "highlight": {
                    "fields": {field: {} for field in self.highlight_fields},
                    "pre_tags": ["<em>"],
                    "post_tags": ["</em>"],
                },
                "size": self.max_results,
            }

            # Execute the search
            response = self.client.search(
                index=self.index_name,
                body=search_query,
            )

            # Process and return the results
            previews = self._process_es_response(response)
            return self._get_full_content(previews)

        except Exception as e:
            safe_msg = self._scrub_error(e)
            logger.warning(f"Error in query_string search: {safe_msg}")
            return []

    def search_by_dsl(self, query_dsl: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        Perform a search using Elasticsearch DSL (Query Domain Specific Language).

        Args:
            query_dsl: The query in Elasticsearch DSL format

        Returns:
            List of search results
        """
        try:
            # Execute the search with the provided DSL
            response = self.client.search(
                index=self.index_name,
                body=query_dsl,
            )

            # Process and return the results
            previews = self._process_es_response(response)
            return self._get_full_content(previews)

        except Exception as e:
            safe_msg = self._scrub_error(e)
            logger.warning(f"Error in DSL search: {safe_msg}")
            return []

    def _process_es_response(self, response: Any) -> List[Dict[str, Any]]:
        """
        Process Elasticsearch response into preview dictionaries.

        Args:
            response: Elasticsearch response dictionary

        Returns:
            List of preview dictionaries
        """
        hits = response.get("hits", {}).get("hits", [])

        # Format results as previews
        previews = []
        for hit in hits:
            source = hit.get("_source", {})
            highlight = hit.get("highlight", {})

            # Extract highlighted snippets or fall back to original content
            snippet = ""
            for field in self.highlight_fields:
                if highlight.get(field):
                    field_snippets = " ... ".join(highlight[field])
                    snippet += field_snippets + " "

            # If no highlights, use a portion of the content
            if not snippet and "content" in source:
                content = source.get("content", "")
                snippet = (
                    content[:SNIPPET_LENGTH_SHORT] + "..."
                    if len(content) > SNIPPET_LENGTH_SHORT
                    else content
                )

            # Create preview object
            preview = {
                "id": hit.get("_id", ""),
                "title": source.get("title", "Untitled Document"),
                "link": source.get("url", "")
                or f"elasticsearch://{self.index_name}/{hit.get('_id', '')}",
                "snippet": snippet.strip(),
                "score": hit.get("_score", 0),
                "_index": hit.get("_index", self.index_name),
            }

            previews.append(preview)

        return previews
