"""
Smart query generation strategy that works for any type of search target.
"""

import concurrent.futures
from typing import Dict, List

from loguru import logger

from ...database.thread_local_session import thread_cleanup
from ..constraints.base_constraint import Constraint
from ..constraints.constraint_analyzer import ConstraintType
from .early_stop_constrained_strategy import EarlyStopConstrainedStrategy


class SmartQueryStrategy(EarlyStopConstrainedStrategy):
    """
    Enhanced strategy with intelligent query generation that:
    1. Analyzes constraints to identify key search terms
    2. Uses LLM to suggest search queries based on constraint meaning
    3. Generates multiple query variations for better coverage
    """

    def __init__(
        self,
        *args,
        use_llm_query_generation: bool = True,
        queries_per_combination: int = 3,
        use_entity_seeding: bool = True,
        use_direct_property_search: bool = True,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.use_llm_query_generation = use_llm_query_generation
        self.queries_per_combination = queries_per_combination
        self.use_entity_seeding = use_entity_seeding
        self.use_direct_property_search = use_direct_property_search
        # Track queries to avoid duplicates
        self.searched_queries = set()
        self.query_variations = set()
        # Store entity seeds for targeted search
        self.entity_seeds = []

    def _build_query(self, constraints: List[Constraint]) -> str:
        """Build intelligent queries using constraint analysis."""
        if self.use_llm_query_generation:
            # Use LLM to generate smart queries
            return self._generate_smart_query(constraints)
        # Fallback to improved standard approach
        return self._build_standard_query(constraints)

    def _generate_smart_query(self, constraints: List[Constraint]) -> str:
        """Use LLM to generate optimal search queries."""
        constraint_text = "\n".join(
            [
                f"- {c.type.value}: {c.value} (weight: {c.weight})"
                for c in constraints
            ]
        )

        # Build a list of already searched queries to avoid duplication
        searched_list = list(self.searched_queries)[:10]  # Show last 10 to LLM
        already_searched = (
            "\n".join([f"- {q}" for q in searched_list])
            if searched_list
            else "None"
        )

        prompt = f"""
Analyze these search constraints and generate an optimal web search query:

Constraints:
{constraint_text}

Target type: {getattr(self, "entity_type", "unknown")}

Already searched queries (avoid these):
{already_searched}

Generate a single search query that would most effectively find results matching these constraints.
The query should:
1. Include the most identifying/unique terms
2. Use appropriate search operators (quotes, AND, OR)
3. Be specific enough to find relevant results but not too narrow
4. Focus on the highest weighted constraints
5. Be different from already searched queries

Return only the search query, nothing else.
"""

        try:
            query = self.model.invoke(prompt).content.strip()

            # Check if this query is too similar to existing ones
            normalized_query = query.strip().lower()
            if normalized_query in self.searched_queries:
                logger.info(
                    f"LLM generated duplicate query, using fallback: {query}"
                )
                return self._build_standard_query(constraints)

            logger.info(f"LLM generated query: {query}")
            return query
        except Exception:
            logger.exception("Failed to generate smart query")
            return self._build_standard_query(constraints)

    def _build_standard_query(self, constraints: List[Constraint]) -> str:
        """Improved standard query building."""
        # Group constraints by importance
        critical_terms = []
        supplementary_terms = []

        for c in constraints:
            term = c.value

            # Quote multi-word terms
            if " " in term and not term.startswith('"'):
                term = f'"{term}"'

            if c.weight > 0.7:
                critical_terms.append(term)
            else:
                supplementary_terms.append(term)

        # Build query with critical terms required, supplementary optional
        query_parts = []

        # Add entity type if known
        entity_type = getattr(self, "entity_type", None)
        if entity_type and entity_type != "unknown entity":
            query_parts.append(entity_type)

        # Add critical terms
        if critical_terms:
            query_parts.extend(critical_terms)

        # Add some supplementary terms
        if supplementary_terms:
            query_parts.extend(
                supplementary_terms[:2]
            )  # Limit to avoid overly specific queries

        return " ".join(query_parts)

    @thread_cleanup
    def _execute_combination_search(self, combo) -> List:
        """Override to generate multiple query variations per combination."""
        all_candidates = []

        if self.use_llm_query_generation:
            # Generate multiple query variations
            queries = self._generate_query_variations(combo.constraints)

            # Execute searches in parallel
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=self.queries_per_combination
            ) as executor:
                futures = []
                for query in queries:
                    # Check if we've already searched this query
                    normalized_query = query.strip().lower()
                    if normalized_query in self.searched_queries:
                        logger.info(f"Skipping duplicate query: '{query}'")
                        continue

                    self.searched_queries.add(normalized_query)
                    future = executor.submit(
                        thread_cleanup(self._execute_search), query
                    )
                    futures.append((query, future))

                for query, future in futures:
                    try:
                        results = future.result()
                        candidates = self._extract_candidates_from_results(
                            results
                        )
                        all_candidates.extend(candidates)

                        logger.info(
                            f"Query '{query}' found {len(candidates)} candidates"
                        )
                    except Exception:
                        logger.exception(f"Search failed for query '{query}'")
        else:
            # Use single query from parent implementation
            candidates = super()._execute_combination_search(combo)
            all_candidates.extend(candidates)

        return all_candidates

    def _generate_query_variations(
        self, constraints: List[Constraint]
    ) -> List[str]:
        """Generate multiple query variations for better coverage."""
        # Handle single constraint case
        if isinstance(constraints, Constraint):
            constraints = [constraints]

        constraint_text = "\n".join(
            [f"- {c.type.value}: {c.value}" for c in constraints]
        )

        # Build a list of already searched queries to avoid duplication
        searched_list = list(self.searched_queries)[:20]  # Show last 20 to LLM
        already_searched = (
            "\n".join([f"- {q}" for q in searched_list])
            if searched_list
            else "None"
        )

        prompt = f"""
Generate {self.queries_per_combination} different search queries for these constraints:

{constraint_text}

Already searched queries (avoid these):
{already_searched}

Each query should:
- Approach the search from a different angle
- Use different search terms or operators
- Target different aspects of the constraints
- Be distinctly different from already searched queries

Provide each query on a separate line.
"""

        try:
            response = self.model.invoke(prompt).content
            queries = [q.strip() for q in response.split("\n") if q.strip()]

            # Filter out duplicates
            unique_queries = []
            for query in queries:
                normalized = query.strip().lower()
                if (
                    normalized not in self.searched_queries
                    and normalized not in self.query_variations
                ):
                    unique_queries.append(query)
                    self.query_variations.add(normalized)
                else:
                    logger.info(
                        f"Filtering out duplicate query variation: {query}"
                    )

            # If all queries were duplicates, generate a fallback
            if not unique_queries:
                fallback = self._build_standard_query(constraints)
                if fallback.strip().lower() not in self.searched_queries:
                    unique_queries = [fallback]

            return unique_queries[: self.queries_per_combination]
        except Exception:
            logger.exception("Failed to generate query variations")
            # Fallback to single query
            return [self._build_standard_query(constraints)]

    def _extract_candidates_from_results(self, results: Dict) -> List:
        """Improved candidate extraction that's more generic."""
        candidates = []
        content = results.get("current_knowledge", "")

        if not content:
            return candidates

        # Use LLM to extract relevant entities/topics from the content
        prompt = f"""
From the following search results, extract all relevant entities, topics, or answers that match our search target type: {getattr(self, "entity_type", "unknown")}

Content:
{content}

List each potential match on a separate line.
Include only names/titles/identifiers, not descriptions.
"""

        try:
            response = self.model.invoke(prompt).content
            entity_names = [
                name.strip() for name in response.split("\n") if name.strip()
            ]

            # Create candidates from extracted names
            from ..candidates.base_candidate import Candidate

            for name in entity_names:
                if name and len(name) < 100:  # Basic validation
                    candidate = Candidate(name=name)
                    candidates.append(candidate)

            logger.info(f"Extracted {len(candidates)} candidates from results")

        except Exception:
            logger.exception("Error extracting candidates")

        return candidates

    def _should_use_entity_seeding(self) -> bool:
        """Determine if entity seeding would be beneficial."""
        entity_type = getattr(self, "entity_type", "").lower()
        return (
            "character" in entity_type
            or "person" in entity_type
            or "hero" in entity_type
        )

    def _perform_entity_seeding(self):
        """Use LLM to suggest specific entity names based on constraints."""
        logger.info("Performing entity seeding based on constraints")

        # Extract key properties from constraints
        key_properties = []
        for constraint in self.constraint_ranking:
            if constraint.weight > 0.7:  # High-weight constraints
                key_properties.append(constraint.value)

        if not key_properties:
            return

        properties_text = "\n".join([f"- {prop}" for prop in key_properties])

        prompt = f"""
Based on these properties, suggest 5-10 specific {self.entity_type} names that might match:

Properties:
{properties_text}

For example, if looking for a scientist from the 19th century, you might suggest:
- Charles Darwin
- Marie Curie
- Louis Pasteur
- Thomas Edison

Provide one name per line. Be specific with actual character/entity names.
"""

        try:
            response = self.model.invoke(prompt).content
            self.entity_seeds = [
                name.strip() for name in response.split("\n") if name.strip()
            ]
            logger.info(f"Generated entity seeds: {self.entity_seeds}")

            # Immediately search for these seeds
            self._search_entity_seeds()

        except Exception:
            logger.exception("Error generating entity seeds")

    def _search_entity_seeds(self):
        """Search for the entity seeds directly."""
        if not self.entity_seeds:
            return

        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
            futures = []
            for seed in self.entity_seeds[:5]:  # Limit to top 5
                query = f'"{seed}"'
                if query.lower() not in self.searched_queries:
                    self.searched_queries.add(query.lower())
                    future = executor.submit(
                        thread_cleanup(self._execute_search), query
                    )
                    futures.append((seed, future))

            for seed, future in futures:
                try:
                    results = future.result()
                    candidates = self._extract_candidates_from_results(results)

                    # Look for exact matches
                    for candidate in candidates:
                        if seed.lower() in candidate.name.lower():
                            logger.info(
                                f"Found seeded entity: {candidate.name}"
                            )
                            # Evaluate immediately
                            if hasattr(self, "_evaluate_candidate_immediately"):
                                self._evaluate_candidate_immediately(candidate)
                            else:
                                # Add to candidates list
                                if not hasattr(self, "candidates"):
                                    self.candidates = []
                                self.candidates.append(candidate)

                except Exception:
                    logger.exception(f"Error searching for seed {seed}")

    def _try_direct_property_search(self):
        """Try direct searches for high-weight property constraints."""
        property_queries = []

        for constraint in self.constraint_ranking:
            if (
                constraint.weight > 0.7
                and constraint.type == ConstraintType.PROPERTY
            ):
                # Create specific property-based queries
                if (
                    "elastic" in constraint.value.lower()
                    or "stretch" in constraint.value.lower()
                ):
                    property_queries.extend(
                        [
                            f'"{constraint.value}" superhero character',
                            f'characters with "{constraint.value}"',
                            f"list of {self.entity_type} {constraint.value}",
                        ]
                    )
                elif (
                    "voice" in constraint.value.lower()
                    or "actor" in constraint.value.lower()
                ):
                    property_queries.append(
                        f"{constraint.value} {self.entity_type}"
                    )

        # Execute property searches
        if property_queries:
            logger.info(
                f"Executing direct property searches: {property_queries}"
            )
            with concurrent.futures.ThreadPoolExecutor(
                max_workers=3
            ) as executor:
                futures = []
                for query in property_queries[
                    :3
                ]:  # Limit to avoid too many searches
                    if query.lower() not in self.searched_queries:
                        self.searched_queries.add(query.lower())
                        future = executor.submit(
                            thread_cleanup(self._execute_search), query
                        )
                        futures.append(future)

                for future in futures:
                    try:
                        results = future.result()
                        candidates = self._extract_candidates_from_results(
                            results
                        )

                        for candidate in candidates:
                            if hasattr(self, "_evaluate_candidate_immediately"):
                                self._evaluate_candidate_immediately(candidate)

                    except Exception:
                        logger.exception("Property search error")

    def _perform_entity_name_search(self):
        """Last resort: search for entity names directly with constraints."""
        logger.info("Performing entity name search fallback")

        for entity_name in self.entity_seeds[:3]:  # Top 3 seeds
            # Combine entity name with key constraints
            constraint_terms = []
            for constraint in self.constraint_ranking[:2]:  # Top 2 constraints
                if constraint.weight > 0.5:
                    constraint_terms.append(constraint.value)

            if constraint_terms:
                query = f'"{entity_name}" {" ".join(constraint_terms)}'
                if query.lower() not in self.searched_queries:
                    logger.info(f"Trying targeted entity search: {query}")
                    self.searched_queries.add(query.lower())

                    try:
                        results = self._execute_search(query)
                        candidates = self._extract_candidates_from_results(
                            results
                        )

                        for candidate in candidates:
                            if entity_name.lower() in candidate.name.lower():
                                logger.info(
                                    f"Found target entity in fallback: {candidate.name}"
                                )
                                if hasattr(
                                    self, "_evaluate_candidate_immediately"
                                ):
                                    self._evaluate_candidate_immediately(
                                        candidate
                                    )

                                    # Check for early stop
                                    if (
                                        hasattr(self, "best_score")
                                        and self.best_score >= 0.99
                                    ):
                                        return

                    except Exception:
                        logger.exception("Entity name search error")

    def _progressive_constraint_search(self):
        """Override to add entity seeding and property search."""
        # Detect entity type first
        self.entity_type = self._detect_entity_type()
        logger.info(f"Detected entity type: {self.entity_type}")

        # Perform entity seeding if enabled and entity type suggests specific entities
        if self.use_entity_seeding and self._should_use_entity_seeding():
            self._perform_entity_seeding()

        # Try direct property search for high-weight properties
        if self.use_direct_property_search:
            self._try_direct_property_search()

        # Continue with normal progressive search
        super()._progressive_constraint_search()

        # If still no good results, try name-based fallback
        if (
            hasattr(self, "best_score")
            and self.best_score < 0.9
            and self.entity_seeds
        ):
            self._perform_entity_name_search()
