"""
Entity-aware source-based search strategy for improved entity identification.
"""

from typing import Dict

from loguru import logger

from ..questions.entity_aware_question import EntityAwareQuestionGenerator
from .source_based_strategy import SourceBasedSearchStrategy


class EntityAwareSourceStrategy(SourceBasedSearchStrategy):
    """
    Enhanced source-based strategy that better handles entity identification queries.

    This strategy:
    1. Detects when queries are seeking specific entities (names, places, etc.)
    2. Generates more targeted search queries for entity identification
    3. Focuses on extracting entity names from search results
    """

    def __init__(
        self,
        search,
        model,
        citation_handler=None,
        include_text_content: bool = True,
        use_cross_engine_filter: bool = True,
        filter_reorder: bool = True,
        filter_reindex: bool = True,
        cross_engine_max_results: int = None,
        all_links_of_system=None,
        use_atomic_facts: bool = False,
    ):
        """Initialize with entity-aware components."""
        # Initialize parent class
        super().__init__(
            search=search,
            model=model,
            citation_handler=citation_handler,
            include_text_content=include_text_content,
            use_cross_engine_filter=use_cross_engine_filter,
            filter_reorder=filter_reorder,
            filter_reindex=filter_reindex,
            cross_engine_max_results=cross_engine_max_results,
            all_links_of_system=all_links_of_system,
            use_atomic_facts=use_atomic_facts,
        )

        # Replace the question generator with entity-aware version
        self.question_generator = EntityAwareQuestionGenerator(self.model)

    def _format_search_results_as_context(self, search_results):
        """Format search results with emphasis on entity extraction."""
        context_snippets = []
        entity_mentions = []

        for i, result in enumerate(search_results):
            title = result.get("title", "Untitled")
            snippet = result.get("snippet", "")
            url = result.get("link", "")

            if snippet:
                context_snippets.append(
                    f"Source {i + 1}: {title}\nURL: {url}\nSnippet: {snippet}"
                )

                # Extract potential entity names (basic approach)
                # In production, this could use NER or more sophisticated extraction
                words = snippet.split()
                for j, word in enumerate(words):
                    # Look for capitalized words that might be names
                    if word and word[0].isupper() and len(word) > 2:
                        # Get some context around the word
                        context_start = max(0, j - 2)
                        context_end = min(len(words), j + 3)
                        entity_context = " ".join(
                            words[context_start:context_end]
                        )
                        if entity_context not in entity_mentions:
                            entity_mentions.append(entity_context)

        # Add entity mentions to context if found
        if entity_mentions:
            context_snippets.append(
                "\nPotential Entity Mentions:\n"
                + "\n".join(f"- {mention}" for mention in entity_mentions[:10])
            )

        return "\n\n".join(context_snippets)

    def analyze_topic(self, query: str) -> Dict:
        """Analyze topic with enhanced entity identification."""
        logger.info(
            f"Starting entity-aware source-based research on topic: {query}"
        )

        # Detect if this is an entity identification query
        entity_keywords = [
            "who",
            "what",
            "which",
            "identify",
            "name",
            "character",
            "person",
            "place",
            "organization",
            "company",
            "author",
            "scientist",
            "inventor",
            "city",
            "country",
            "book",
            "movie",
        ]

        is_entity_query = any(
            keyword in query.lower() for keyword in entity_keywords
        )

        if is_entity_query:
            logger.info(
                "Detected entity identification query - using enhanced search patterns"
            )
            self._update_progress(
                "Detected entity identification query - optimizing search strategy",
                3,
                {
                    "phase": "init",
                    "query_type": "entity_identification",
                    "strategy": "entity_aware_source",
                },
            )

        # Use parent's analyze_topic with our enhanced components
        return super().analyze_topic(query)
