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

from elasticsearch import Elasticsearch
from langchain_core.language_models import BaseLLM

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


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

    def __init__(
        self,
        hosts: List[str] = ["http://localhost:9200"],
        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,
    ):
        """
        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,
        )

        self.index_name = index_name
        self.highlight_fields = highlight_fields
        self.search_fields = search_fields
        self.filter_query = filter_query or {}

        # Initialize the Elasticsearch client
        es_args = {}

        # 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:
            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:
            logger.exception(f"Failed to connect to Elasticsearch: {e!s}")
            raise ConnectionError(f"Could not connect to Elasticsearch: {e!s}")

    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[:250] + "..." if len(content) > 250 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:
            logger.exception(f"Error getting Elasticsearch previews: {e!s}")
            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
        """
        # 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 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:
                logger.exception(
                    f"Error fetching full content for document {doc_id}: {e!s}"
                )
                # 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:
            logger.exception(f"Error in query_string search: {e!s}")
            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:
            logger.exception(f"Error in DSL search: {e!s}")
            return []

    def _process_es_response(
        self, response: Dict[str, 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[:250] + "..." if len(content) > 250 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
