"""
BrowseComp-Optimized Search Strategy for Complex Query Solving

This strategy is specifically designed to handle BrowseComp-style puzzle queries
where specific clues need to be matched to find a location, person, or event.
"""

from dataclasses import dataclass, field
from datetime import datetime, UTC
from typing import Any, Dict, List, Optional

from langchain_core.language_models import BaseChatModel
from loguru import logger

from ...utilities.search_utilities import format_findings, remove_think_tags
from ..findings.repository import FindingsRepository
from .base_strategy import BaseSearchStrategy
from .source_based_strategy import SourceBasedSearchStrategy


@dataclass
class QueryClues:
    """Extract specific clues from a complex query."""

    location_clues: List[str] = field(default_factory=list)
    temporal_clues: List[str] = field(default_factory=list)
    numerical_clues: List[str] = field(default_factory=list)
    name_clues: List[str] = field(default_factory=list)
    incident_clues: List[str] = field(default_factory=list)
    comparison_clues: List[str] = field(default_factory=list)
    all_clues: List[str] = field(default_factory=list)
    query_type: str = "unknown"  # location, person, event, object, etc.


class BrowseCompOptimizedStrategy(BaseSearchStrategy):
    """
    A strategy optimized for solving BrowseComp-style puzzle queries.

    Key features:
    1. Extracts specific clues from the query
    2. Searches for combinations of clues
    3. Progressively narrows down candidates
    4. Uses specific place/person/object names when found
    5. Verifies candidates against all clues
    """

    def __init__(
        self,
        model: BaseChatModel,
        search: Any,
        all_links_of_system: List[Any],
        max_browsecomp_iterations: int = 15,
        confidence_threshold: float = 0.90,
        max_iterations: int = 1,  # This is for source-based strategy iterations
        questions_per_iteration: int = 3,  # This is for source-based strategy questions
        settings_snapshot=None,
    ):
        """Initialize the BrowseComp-optimized strategy."""
        super().__init__(all_links_of_system, settings_snapshot)
        self.model = model
        self.search = search
        self.max_browsecomp_iterations = max_browsecomp_iterations
        self.confidence_threshold = confidence_threshold
        self.findings_repository = FindingsRepository(model)

        # Store the source-based strategy parameters
        self.source_max_iterations = max_iterations
        self.source_questions_per_iteration = questions_per_iteration

        # State tracking
        self.query_clues: Optional[QueryClues] = None
        self.confirmed_info: Dict[str, Any] = {}
        self.candidates: List[Dict[str, Any]] = []
        self.search_history: List[str] = []
        self.iteration: int = 0

    def analyze_topic(self, query: str) -> Dict:
        """Analyze a topic using BrowseComp-optimized approach."""
        # Initialize
        self.all_links_of_system.clear()
        self.questions_by_iteration = []
        self.findings = []
        self.iteration = 0

        # Extract clues from query
        self.query_clues = self._extract_clues(query)

        # Progress callback
        if self.progress_callback:
            self.progress_callback(
                f"Identified {len(self.query_clues.all_clues)} clues from query",
                1,
                {
                    "phase": "init",
                    "strategy": "browsecomp_optimized",
                    "clues_count": len(self.query_clues.all_clues),
                    "query_type": self.query_clues.query_type,
                },
            )

        logger.info(f"Starting BrowseComp optimization for: {query}")
        logger.info(f"Extracted {len(self.query_clues.all_clues)} clues")

        # Add initial analysis
        initial_finding = {
            "phase": "Initial Analysis",
            "content": f"""
**Query**: {query}

**Strategy**: BrowseComp-Optimized
- Query type: {self.query_clues.query_type}
- Total clues: {len(self.query_clues.all_clues)}
- Location clues: {len(self.query_clues.location_clues)}
- Temporal clues: {len(self.query_clues.temporal_clues)}
- Numerical clues: {len(self.query_clues.numerical_clues)}

**Key Clues**:
{chr(10).join(f"- {clue}" for clue in self.query_clues.all_clues[:5])}

**Starting systematic search**...
""".strip(),
            "timestamp": self._get_timestamp(),
        }
        self.findings.append(initial_finding)

        # Main search loop
        while self.iteration < self.max_browsecomp_iterations:
            self.iteration += 1

            # Progress update
            if self.progress_callback:
                progress = (
                    int((self.iteration / self.max_browsecomp_iterations) * 80)
                    + 10
                )
                self.progress_callback(
                    f"Iteration {self.iteration}: {len(self.candidates)} candidates, {len(self.confirmed_info)} confirmed facts",
                    progress,
                    {
                        "phase": "searching",
                        "iteration": self.iteration,
                        "candidates_count": len(self.candidates),
                        "confirmed_facts": len(self.confirmed_info),
                    },
                )

            # Generate targeted search query
            search_query = self._generate_targeted_search()

            if not search_query:
                logger.info("No more searches needed - sufficient candidates")
                break

            # Execute search
            logger.info(
                f"Iteration {self.iteration}: Searching for '{search_query}'"
            )
            search_results = self._execute_search(search_query)

            # Process results
            self._process_search_results(search_results, search_query)

            # Check if we have a strong candidate
            if self._evaluate_candidates():
                break

            # Add iteration finding
            iteration_summary = f"""
**Search Query**: {search_query}

**Candidates Found**: {len(self.candidates)}
{chr(10).join(f"- {c['name']} (confidence: {c['confidence']:.0%})" for c in self.candidates[:3])}

**Confirmed Facts**: {len(self.confirmed_info)}
{chr(10).join(f"- {k}: {v}" for k, v in list(self.confirmed_info.items())[:3])}

**Progress**: Iteration {self.iteration}/{self.max_browsecomp_iterations}
"""
            self.findings.append(
                {
                    "phase": f"Iteration {self.iteration}",
                    "content": iteration_summary.strip(),
                    "timestamp": self._get_timestamp(),
                }
            )

        # Generate final answer
        final_result = self._synthesize_final_answer(query)

        if self.progress_callback:
            self.progress_callback(
                f"Analysis complete - {self.iteration} iterations, {len(self.candidates)} final candidates",
                100,
                {
                    "phase": "complete",
                    "strategy": "browsecomp_optimized",
                    "total_iterations": self.iteration,
                    "final_candidates": len(self.candidates),
                },
            )

        return final_result

    def _extract_clues(self, query: str) -> QueryClues:
        """Extract specific clues from the query."""
        prompt = f"""
Analyze this query and extract ALL specific clues that help identify the answer.

Query: {query}

Extract the following types of clues (BE VERY SPECIFIC AND COMPREHENSIVE):
1. Location clues (geographical features, regions, landmarks)
2. Temporal clues (dates, time periods, years - extract exact years/ranges)
3. Numerical clues (statistics, counts, comparisons - extract exact numbers)
4. Name clues (hints about the name, body parts, colors, etc.)
5. Incident clues (accidents, events, activities - be specific about what happened)
6. Comparison clues (comparisons to other places/things - extract exact comparison ratios)

IMPORTANT: Extract EVERY specific detail, number, date, and comparison mentioned.

Also determine the query type: location, person, event, object, or other.

Format your response as:
QUERY_TYPE: [type]

LOCATION_CLUES:
- [clue 1]
- [clue 2]

TEMPORAL_CLUES:
- [clue 1 with exact dates/years]

NUMERICAL_CLUES:
- [clue 1 with exact numbers]

NAME_CLUES:
- [clue 1]

INCIDENT_CLUES:
- [clue 1 with specific details]

COMPARISON_CLUES:
- [clue 1 with exact comparison ratio]

ALL_CLUES_SUMMARY:
- [most important clue 1]
- [most important clue 2]
- [most important clue 3]
- [most important clue 4]
- [most important clue 5]
"""

        response = self.model.invoke(prompt)
        content = remove_think_tags(str(response.content))

        clues = QueryClues()
        current_section = None

        for line in content.strip().split("\n"):
            line = line.strip()

            if line.startswith("QUERY_TYPE:"):
                clues.query_type = line.split(":", 1)[1].strip().lower()
            elif line.endswith("_CLUES:") or line == "ALL_CLUES_SUMMARY:":
                current_section = (
                    line.replace("_CLUES:", "")
                    .replace("ALL_CLUES_SUMMARY:", "all")
                    .lower()
                )
            elif line.startswith("-") and current_section:
                clue = line[1:].strip()
                if current_section == "location":
                    clues.location_clues.append(clue)
                elif current_section == "temporal":
                    clues.temporal_clues.append(clue)
                elif current_section == "numerical":
                    clues.numerical_clues.append(clue)
                elif current_section == "name":
                    clues.name_clues.append(clue)
                elif current_section == "incident":
                    clues.incident_clues.append(clue)
                elif current_section == "comparison":
                    clues.comparison_clues.append(clue)
                elif current_section == "all":
                    clues.all_clues.append(clue)

        # Ensure we have all clues
        if not clues.all_clues:
            clues.all_clues = (
                clues.location_clues
                + clues.temporal_clues
                + clues.numerical_clues
                + clues.name_clues
                + clues.incident_clues
                + clues.comparison_clues
            )[:5]  # Top 5 most important

        return clues

    def _generate_targeted_search(self) -> Optional[str]:
        """Generate a targeted search query based on current knowledge."""
        # If we have specific candidates, search for verification
        if self.candidates and self.iteration > 2:
            top_candidate = self.candidates[0]

            # Search for specific verification of top candidate with unverified clues
            unverified_clues = self._get_unverified_clues(top_candidate)

            if unverified_clues:
                # Pick the most specific unverified clue
                if any(
                    "fell" in clue or "accident" in clue
                    for clue in unverified_clues
                ):
                    return f"{top_candidate['name']} accident fall death {' '.join(self.query_clues.temporal_clues[:1])}"
                if any(
                    "search and rescue" in clue.lower() or "sar" in clue.lower()
                    for clue in unverified_clues
                ):
                    return f"{top_candidate['name']} search and rescue incidents 2014 statistics"
                if any("84.5" in clue for clue in unverified_clues):
                    return f"{top_candidate['name']} 2014 search rescue statistics Grand Canyon 2022 comparison"

        # Initial searches - combine multiple clues
        if self.iteration <= 3:
            # First iteration - broad search with key clues
            if self.iteration == 1:
                key_terms = []
                if self.query_clues.location_clues:
                    key_terms.extend(self.query_clues.location_clues[:1])
                if self.query_clues.name_clues:
                    key_terms.extend(self.query_clues.name_clues[:1])
                if self.query_clues.query_type == "location":
                    key_terms.append("hiking trail scenic viewpoint")
                return " ".join(key_terms)

            # Second iteration - add temporal/incident info
            if self.iteration == 2:
                key_terms = []
                if self.query_clues.temporal_clues:
                    key_terms.extend(self.query_clues.temporal_clues[:1])
                if self.query_clues.incident_clues:
                    key_terms.extend(self.query_clues.incident_clues[:1])
                if self.query_clues.location_clues:
                    key_terms.extend(self.query_clues.location_clues[:1])
                return " ".join(key_terms)

        # Middle iterations - search for specific combinations
        elif 3 <= self.iteration <= 8:
            # Try different clue combinations
            combinations = [
                (self.query_clues.location_clues, self.query_clues.name_clues),
                (
                    self.query_clues.temporal_clues,
                    self.query_clues.incident_clues,
                ),
                (
                    self.query_clues.numerical_clues,
                    self.query_clues.location_clues,
                ),
                (self.query_clues.name_clues, self.query_clues.incident_clues),
            ]

            combo_idx = (self.iteration - 3) % len(combinations)
            clues1, clues2 = combinations[combo_idx]

            terms = []
            if clues1:
                terms.extend(clues1[:1])
            if clues2:
                terms.extend(clues2[:1])

            return " ".join(terms)

        # Late iterations - search for specific statistics or comparisons
        elif self.iteration > 8:
            if self.query_clues.comparison_clues:
                return " ".join(
                    self.query_clues.comparison_clues[:1] + ["statistics data"]
                )
            if self.query_clues.numerical_clues:
                return " ".join(
                    self.query_clues.numerical_clues[:1]
                    + ["official statistics 2014 2022"]
                )

        # Default - use remaining clues
        all_unused = [
            c
            for c in self.query_clues.all_clues
            if c not in self.search_history
        ]
        if all_unused:
            return all_unused[0]

        return None

    def _execute_search(self, search_query: str) -> Dict:
        """Execute a search using source-based strategy."""
        # Track search history
        self.search_history.append(search_query)

        # Use source-based strategy
        source_strategy = SourceBasedSearchStrategy(
            model=self.model,
            search=self.search,
            all_links_of_system=self.all_links_of_system,
            include_text_content=True,
            use_cross_engine_filter=True,
            use_atomic_facts=True,
        )

        source_strategy.max_iterations = self.source_max_iterations
        source_strategy.questions_per_iteration = (
            self.source_questions_per_iteration
        )

        if self.progress_callback:

            def wrapped_callback(message, progress, data):
                data["parent_iteration"] = self.iteration
                data["parent_strategy"] = "browsecomp_optimized"
                self.progress_callback(
                    f"Iteration {self.iteration}: {message}", progress, data
                )

            source_strategy.set_progress_callback(wrapped_callback)

        results = source_strategy.analyze_topic(search_query)

        if "questions_by_iteration" in results:
            self.questions_by_iteration.extend(
                results["questions_by_iteration"]
            )

        return results

    def _process_search_results(self, search_results: Dict, search_query: str):
        """Process search results and update candidates."""
        current_knowledge = search_results.get("current_knowledge", "")

        if not current_knowledge:
            return

        prompt = f"""
Based on the search results, extract specific information about potential answers.

Query Clues:
- Query type: {self.query_clues.query_type}
- Key clues: {", ".join(self.query_clues.all_clues[:3])}

Search Query: {search_query}

Search Results:
{current_knowledge[:3000]}

Current Candidates: {len(self.candidates)}
{chr(10).join(f"- {c['name']} ({c['confidence']:.0%})" for c in self.candidates[:3])}

Extract:
1. SPECIFIC_NAMES: Any specific place names, trail names, or landmarks mentioned
2. CONFIRMED_FACTS: Facts that match our clues (with which clue they match)
3. NEW_CANDIDATES: New potential answers with confidence (0-1)
4. ELIMINATED_CANDIDATES: Candidates we can rule out

Format:
SPECIFIC_NAMES:
- [name 1]
- [name 2]

CONFIRMED_FACTS:
- [fact]: matches [clue]

NEW_CANDIDATES:
- [name]: [confidence] (reason)

ELIMINATED_CANDIDATES:
- [name]: [reason for elimination]
"""

        response = self.model.invoke(prompt)
        content = remove_think_tags(str(response.content))

        # Parse response
        current_section = None

        for line in content.strip().split("\n"):
            line = line.strip()

            if line.endswith(":"):
                current_section = line[:-1].lower()
            elif line.startswith("-") and current_section:
                item = line[1:].strip()

                if current_section == "specific_names":
                    # Add as potential candidate if not already present
                    if not any(
                        c["name"].lower() == item.lower()
                        for c in self.candidates
                    ):
                        self.candidates.append(
                            {
                                "name": item,
                                "confidence": 0.7,
                                "source": search_query,
                                "matched_clues": [],
                            }
                        )

                elif current_section == "confirmed_facts":
                    if ":" in item:
                        fact, clue = item.split(":", 1)
                        self.confirmed_info[fact.strip()] = clue.strip()

                elif current_section == "new_candidates":
                    if ":" in item:
                        parts = item.split(":", 1)
                        name = parts[0].strip()

                        # Parse confidence
                        confidence_str = parts[1].strip()
                        try:
                            confidence = float(confidence_str.split()[0])
                        except (ValueError, IndexError):
                            confidence = 0.5

                        # Update or add candidate
                        existing = False
                        for candidate in self.candidates:
                            if candidate["name"].lower() == name.lower():
                                candidate["confidence"] = max(
                                    candidate["confidence"], confidence
                                )
                                existing = True
                                break

                        if not existing:
                            self.candidates.append(
                                {
                                    "name": name,
                                    "confidence": confidence,
                                    "source": search_query,
                                    "matched_clues": [],
                                }
                            )

                elif current_section == "eliminated_candidates":
                    if ":" in item:
                        name = item.split(":")[0].strip()
                        self.candidates = [
                            c
                            for c in self.candidates
                            if c["name"].lower() != name.lower()
                        ]

        # Update matched clues for candidates
        for candidate in self.candidates:
            candidate["matched_clues"] = self._get_matched_clues(
                candidate["name"]
            )
            # Update confidence based on matched clues
            clue_confidence = len(candidate["matched_clues"]) / len(
                self.query_clues.all_clues
            )
            candidate["confidence"] = (
                candidate["confidence"] + clue_confidence
            ) / 2

        # Sort candidates by confidence
        self.candidates.sort(key=lambda x: x["confidence"], reverse=True)

        # Keep only top candidates
        self.candidates = self.candidates[:5]

    def _get_matched_clues(self, candidate_name: str) -> List[str]:
        """Determine which clues a candidate matches."""
        matched = []

        # Check name clues
        for clue in self.query_clues.name_clues:
            if any(
                word.lower() in candidate_name.lower() for word in clue.split()
            ):
                matched.append(f"name: {clue}")

        # Check location clues
        for clue in self.query_clues.location_clues:
            if (
                candidate_name.lower()
                in self.confirmed_info.get(clue, "").lower()
            ):
                matched.append(f"location: {clue}")

        # Check confirmed facts
        for fact, clue in self.confirmed_info.items():
            if candidate_name.lower() in fact.lower():
                matched.append(f"fact: {clue}")

        # Check temporal clues if mentioned in confirmed facts
        for clue in self.query_clues.temporal_clues:
            for fact in self.confirmed_info:
                if candidate_name.lower() in fact.lower() and any(
                    year in fact for year in clue.split() if year.isdigit()
                ):
                    matched.append(f"temporal: {clue}")

        # Check incident clues
        for clue in self.query_clues.incident_clues:
            for fact in self.confirmed_info:
                if candidate_name.lower() in fact.lower() and any(
                    word in fact.lower() for word in clue.lower().split()
                ):
                    matched.append(f"incident: {clue}")

        return list(set(matched))  # Remove duplicates

    def _get_unverified_clues(self, candidate: Dict) -> List[str]:
        """Get clues that haven't been verified for a candidate."""
        all_clues = (
            self.query_clues.location_clues
            + self.query_clues.temporal_clues
            + self.query_clues.numerical_clues
            + self.query_clues.name_clues
            + self.query_clues.incident_clues
            + self.query_clues.comparison_clues
        )

        matched_clues = candidate.get("matched_clues", [])
        matched_clue_texts = [
            clue.split(": ", 1)[1] if ": " in clue else clue
            for clue in matched_clues
        ]

        unverified = []
        for clue in all_clues:
            if not any(
                clue in matched_text for matched_text in matched_clue_texts
            ):
                unverified.append(clue)

        return unverified

    def _evaluate_candidates(self) -> bool:
        """Check if we have a sufficiently confident answer."""
        if not self.candidates:
            return False

        top_candidate = self.candidates[0]

        # Need high confidence and clue matching
        if top_candidate["confidence"] >= self.confidence_threshold:
            matched_ratio = len(top_candidate["matched_clues"]) / len(
                self.query_clues.all_clues
            )
            # For BrowseComp, require higher clue matching threshold
            if matched_ratio >= 0.8:  # At least 80% of clues matched
                return True

        # Very high confidence only if ALL clues matched
        if top_candidate["confidence"] >= 0.95:
            matched_ratio = len(top_candidate["matched_clues"]) / len(
                self.query_clues.all_clues
            )
            if matched_ratio >= 0.9:  # At least 90% clues matched
                return True

        # For BrowseComp queries, always do at least 3 iterations
        if self.iteration < 3:
            return False

        # Or if we've done enough iterations and have a clear leader
        if self.iteration >= 8:
            if len(self.candidates) > 1:
                confidence_gap = (
                    top_candidate["confidence"]
                    - self.candidates[1]["confidence"]
                )
                if confidence_gap > 0.3:  # 30% gap to second place
                    return True

        return False

    def _synthesize_final_answer(self, original_query: str) -> Dict:
        """Generate final answer in BrowseComp format."""
        if self.candidates:
            top_candidate = self.candidates[0]
            answer = top_candidate["name"]
            confidence = int(top_candidate["confidence"] * 100)
        else:
            answer = "Unable to determine"
            confidence = 0

        # Generate explanation
        prompt = f"""
Provide a concise final answer to this query:
{original_query}

Based on our research:
- Top answer: {answer}
- Confidence: {confidence}%
- Matched clues: {len(top_candidate["matched_clues"]) if self.candidates else 0}/{len(self.query_clues.all_clues)}

Confirmed facts:
{chr(10).join(f"- {k}: {v}" for k, v in self.confirmed_info.items())}

Format your response EXACTLY as:
Explanation: {{brief explanation of why this answer matches the clues}}
Exact Answer: {{the exact answer - just the name/place/etc}}
Confidence: {{confidence}}%
"""

        response = self.model.invoke(prompt)
        final_answer = remove_think_tags(str(response.content))

        # Add comprehensive findings
        synthesis_summary = f"""
**Final Answer**: {answer}

**Research Summary**:
- Completed {self.iteration} search iterations
- Evaluated {len(self.candidates)} candidates
- Matched {len(top_candidate["matched_clues"]) if self.candidates else 0}/{len(self.query_clues.all_clues)} clues
- Final confidence: {confidence}%

**Top Candidates**:
{chr(10).join(f"{i + 1}. {c['name']} ({c['confidence']:.0%})" for i, c in enumerate(self.candidates[:3]))}

**Confirmed Facts**:
{chr(10).join(f"- {k}: {v}" for k, v in self.confirmed_info.items())}

**Search History**:
{chr(10).join(f"{i + 1}. {q}" for i, q in enumerate(self.search_history))}
"""

        self.findings.append(
            {
                "phase": "Final Synthesis",
                "content": synthesis_summary,
                "timestamp": self._get_timestamp(),
            }
        )

        # Compile questions
        questions_dict = {}
        for i, questions in enumerate(self.questions_by_iteration):
            if isinstance(questions, list):
                questions_dict[i + 1] = questions
            elif isinstance(questions, dict):
                questions_dict.update(questions)

        # Format findings
        formatted_findings = format_findings(
            self.findings, final_answer, questions_dict
        )

        return {
            "current_knowledge": final_answer,
            "formatted_findings": formatted_findings,
            "findings": self.findings,
            "iterations": self.iteration,
            "questions_by_iteration": questions_dict,
            "all_links_of_system": self.all_links_of_system,
            "sources": self.all_links_of_system,
            "candidates": self.candidates,
            "confirmed_info": self.confirmed_info,
            "clues": {
                "location": self.query_clues.location_clues,
                "temporal": self.query_clues.temporal_clues,
                "numerical": self.query_clues.numerical_clues,
                "name": self.query_clues.name_clues,
                "incident": self.query_clues.incident_clues,
                "comparison": self.query_clues.comparison_clues,
            },
            "strategy": "browsecomp_optimized",
        }

    def _get_timestamp(self) -> str:
        """Get current timestamp for findings."""
        return datetime.now(UTC).isoformat()
