"""
Progressive explorer for BrowseComp-style systematic search exploration.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Set, Tuple

from loguru import logger

from ...utilities.thread_context import preserve_research_context
from ..parallel_search import run_parallel_searches


@dataclass
class SearchProgress:
    """Track search progress and findings."""

    searched_terms: Set[str] = field(default_factory=set)
    found_candidates: Dict[str, float] = field(
        default_factory=dict
    )  # name -> confidence
    verified_facts: Dict[str, str] = field(
        default_factory=dict
    )  # fact -> source
    entity_coverage: Dict[str, Set[str]] = field(
        default_factory=dict
    )  # entity_type -> searched_entities
    search_depth: int = 0

    def update_coverage(self, entity_type: str, entity: str):
        """Update entity coverage tracking."""
        if entity_type not in self.entity_coverage:
            self.entity_coverage[entity_type] = set()
        self.entity_coverage[entity_type].add(entity.lower())

    def get_uncovered_entities(
        self, entities: Dict[str, List[str]]
    ) -> Dict[str, List[str]]:
        """Get entities that haven't been searched yet."""
        uncovered = {}
        for entity_type, entity_list in entities.items():
            covered = self.entity_coverage.get(entity_type, set())
            uncovered_list = [
                e for e in entity_list if e.lower() not in covered
            ]
            if uncovered_list:
                uncovered[entity_type] = uncovered_list
        return uncovered


class ProgressiveExplorer:
    """
    Explorer that implements progressive search strategies for BrowseComp.

    Key features:
    1. Tracks search progress to avoid redundancy
    2. Progressively combines entities
    3. Identifies and pursues promising candidates
    4. Maintains simple approach without over-filtering
    """

    def __init__(self, search_engine, model):
        self.search_engine = search_engine
        self.model = model
        self.progress = SearchProgress()
        self.max_results_per_search = 20  # Keep more results

    def explore(
        self,
        queries: List[str],
        constraints: List = None,
        max_workers: int = 5,
        extracted_entities: Dict[str, List[str]] = None,
    ) -> Tuple[List, SearchProgress]:
        """
        Execute progressive exploration with entity tracking.

        Returns both candidates and search progress for strategy use.
        """
        all_results = []
        extracted_entities = extracted_entities or {}

        # Execute searches in parallel (like source-based strategy)
        search_results = self._parallel_search(queries, max_workers)

        # Process results without filtering (trust the LLM later)
        for query, results in search_results:
            self.progress.searched_terms.add(query.lower())

            # Track which entities were covered in this search
            self._update_entity_coverage(query, extracted_entities)

            # Extract any specific names/candidates from results
            candidates = self._extract_candidates_from_results(results, query)
            for candidate_name, confidence in candidates.items():
                if candidate_name in self.progress.found_candidates:
                    # Update confidence if higher
                    self.progress.found_candidates[candidate_name] = max(
                        self.progress.found_candidates[candidate_name],
                        confidence,
                    )
                else:
                    self.progress.found_candidates[candidate_name] = confidence

            # Keep all results for final synthesis
            all_results.extend(results)

        self.progress.search_depth += 1

        # Return both results and progress
        return all_results, self.progress

    def generate_verification_searches(
        self,
        candidates: Dict[str, float],
        constraints: List,
        max_searches: int = 5,
    ) -> List[str]:
        """Generate targeted searches to verify top candidates."""
        if not candidates:
            return []

        # Get top candidates by confidence
        top_candidates = sorted(
            candidates.items(), key=lambda x: x[1], reverse=True
        )[:3]

        verification_searches = []
        for candidate_name, confidence in top_candidates:
            # Generate verification searches for this candidate
            for constraint in constraints[:2]:  # Verify top constraints
                search = f'"{candidate_name}" {constraint.description}'
                if search.lower() not in self.progress.searched_terms:
                    verification_searches.append(search)

        return verification_searches[:max_searches]

    def _extract_candidates_from_results(
        self, results: List[Dict], query: str
    ) -> Dict[str, float]:
        """Extract potential answer candidates from search results."""
        candidates = {}

        # Simple extraction based on titles and snippets
        for result in results[:10]:  # Focus on top results
            title = result.get("title", "")
            snippet = result.get("snippet", "")

            # Look for proper nouns and specific names
            # This is simplified - in practice, might use NER or more sophisticated extraction
            combined_text = f"{title} {snippet}"

            # Extract quoted terms as potential candidates
            import re

            quoted_terms = re.findall(r'"([^"]+)"', combined_text)
            for term in quoted_terms:
                if (
                    len(term) > 2 and len(term) < 50
                ):  # Reasonable length for an answer
                    candidates[term] = 0.3  # Base confidence from appearance

            # Boost confidence if appears in title
            if title:
                # Titles often contain the actual answer
                title_words = title.split()
                for i in range(len(title_words)):
                    for j in range(i + 1, min(i + 4, len(title_words) + 1)):
                        phrase = " ".join(title_words[i:j])
                        if (
                            len(phrase) > 3 and phrase[0].isupper()
                        ):  # Likely proper noun
                            candidates[phrase] = candidates.get(phrase, 0) + 0.2

        return candidates

    def _update_entity_coverage(
        self, query: str, entities: Dict[str, List[str]]
    ):
        """Track which entities have been covered in searches."""
        query_lower = query.lower()

        for entity_type, entity_list in entities.items():
            for entity in entity_list:
                if entity.lower() in query_lower:
                    self.progress.update_coverage(entity_type, entity)

    def suggest_next_searches(
        self, entities: Dict[str, List[str]], max_suggestions: int = 5
    ) -> List[str]:
        """Suggest next searches based on coverage and findings."""
        suggestions = []

        # 1. Check uncovered entities
        uncovered = self.progress.get_uncovered_entities(entities)

        # 2. If we have candidates, verify them with uncovered constraints
        if self.progress.found_candidates:
            top_candidate = max(
                self.progress.found_candidates.items(), key=lambda x: x[1]
            )[0]

            # Combine candidate with uncovered entities
            for entity_list in uncovered.values():
                for entity in entity_list[:2]:
                    search = f'"{top_candidate}" {entity}'
                    if search.lower() not in self.progress.searched_terms:
                        suggestions.append(search)

        # 3. Otherwise, create new combinations of uncovered entities
        else:
            # Focus on systematic coverage
            if uncovered.get("temporal"):
                # Year-by-year with key term
                key_term = (
                    entities.get("names", [""])[0]
                    or entities.get("descriptors", [""])[0]
                )
                for year in uncovered["temporal"][:3]:
                    search = f"{key_term} {year}".strip()
                    if search.lower() not in self.progress.searched_terms:
                        suggestions.append(search)

            if uncovered.get("names") and uncovered.get("descriptors"):
                # Combine names with descriptors
                for name in uncovered["names"][:2]:
                    for desc in uncovered["descriptors"][:2]:
                        search = f"{name} {desc}"
                        if search.lower() not in self.progress.searched_terms:
                            suggestions.append(search)

        return suggestions[:max_suggestions]

    def _parallel_search(
        self, queries: List[str], max_workers: int
    ) -> List[Tuple[str, List[Dict]]]:
        """Execute searches in parallel and return results."""
        results = []

        def search_query(query):
            try:
                search_results = self.search_engine.run(query)
                return search_results or []
            except Exception:
                logger.exception(f"Error searching '{query}'")
                return []

        # Create context-preserving wrapper for the search function
        context_aware_search = preserve_research_context(search_query)

        # Run searches in parallel
        completed = run_parallel_searches(
            queries, context_aware_search, max_workers=max_workers
        )
        results.extend(completed)

        return results
