"""
Early-stop constrained search strategy that evaluates candidates immediately
and stops when finding a very high confidence match.
"""

import concurrent.futures
import threading
from typing import Dict, List

from loguru import logger

from ...database.thread_local_session import thread_cleanup
from ..candidates.base_candidate import Candidate
from ..constraints.base_constraint import Constraint
from .parallel_constrained_strategy import ParallelConstrainedStrategy


class EarlyStopConstrainedStrategy(ParallelConstrainedStrategy):
    """
    Enhanced constrained strategy that:
    1. Evaluates candidates as soon as they're found
    2. Stops early when finding a very high confidence match (99%+)
    3. Runs evaluation and search concurrently
    """

    def __init__(
        self,
        *args,
        early_stop_threshold: float = 0.99,
        concurrent_evaluation: bool = True,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.early_stop_threshold = early_stop_threshold
        self.concurrent_evaluation = concurrent_evaluation

        # Thread-safe tracking
        self.found_answer = threading.Event()
        self.best_candidate = None
        self.best_score = 0.0
        self.evaluation_lock = threading.Lock()

        # Track candidates being evaluated
        self.evaluating_candidates = set()
        self.evaluated_candidates = {}

    def _parallel_search(self, combinations: List) -> List[Candidate]:
        """Execute searches in parallel with immediate candidate evaluation."""
        all_candidates = []
        evaluation_futures = []

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=self.parallel_workers
        ) as executor:
            # Submit all searches
            search_futures = {
                executor.submit(self._execute_combination_search, combo): combo
                for combo in combinations
            }

            # Process results as they complete
            for future in concurrent.futures.as_completed(search_futures):
                # Check if we should stop early
                if self.found_answer.is_set():
                    logger.info(
                        f"Early stop triggered - found answer: {self.best_candidate}"
                    )
                    break

                combo = search_futures[future]
                try:
                    candidates = future.result()
                    all_candidates.extend(candidates)

                    # Start evaluating candidates immediately if concurrent evaluation is enabled
                    if self.concurrent_evaluation:
                        for candidate in candidates:
                            if candidate.name not in self.evaluating_candidates:
                                self.evaluating_candidates.add(candidate.name)
                                eval_future = executor.submit(
                                    thread_cleanup(
                                        self._evaluate_candidate_immediately
                                    ),
                                    candidate,
                                )
                                evaluation_futures.append(eval_future)

                    if self.progress_callback:
                        self.progress_callback(
                            f"Found {len(candidates)} candidates, evaluating...",
                            None,
                            {
                                "phase": "parallel_search_with_eval",
                                "candidates": len(all_candidates),
                                "best_score": self.best_score,
                                "best_candidate": self.best_candidate,
                            },
                        )

                except Exception:
                    logger.exception(f"Search failed for {combo.query}")

            # Wait for evaluation futures to complete
            for future in concurrent.futures.as_completed(evaluation_futures):
                try:
                    future.result()
                except Exception:
                    logger.exception("Evaluation failed")

        return all_candidates

    def _evaluate_candidate_immediately(self, candidate: Candidate) -> float:
        """Evaluate a candidate against all constraints immediately."""
        try:
            logger.info(f"Immediately evaluating candidate: {candidate.name}")

            # Calculate overall score across all constraints
            total_score = 0.0
            constraint_scores = []

            for constraint in self.constraint_ranking:
                # Get evidence for this constraint
                evidence = self._gather_evidence_for_constraint(
                    candidate, constraint
                )
                score = self._evaluate_evidence(evidence, constraint)
                constraint_scores.append(score)

                # Update progress
                if self.progress_callback:
                    symbol = "✓" if score >= 0.8 else "○"
                    self.progress_callback(
                        f"{symbol} {candidate.name} | {constraint.type.value}: {int(score * 100)}%",
                        None,
                        {
                            "phase": "immediate_evaluation",
                            "candidate": candidate.name,
                            "constraint": constraint.value,
                            "score": score,
                        },
                    )

                # If this candidate fails a critical constraint badly, skip remaining checks
                if score < 0.3 and constraint.weight > 0.8:
                    logger.info(
                        f"Candidate {candidate.name} failed critical constraint early"
                    )
                    break

            # Calculate average score
            if constraint_scores:
                total_score = sum(constraint_scores) / len(constraint_scores)

            # Thread-safe update of best candidate
            with self.evaluation_lock:
                self.evaluated_candidates[candidate.name] = total_score

                if total_score > self.best_score:
                    self.best_score = total_score
                    self.best_candidate = candidate.name

                    logger.info(
                        f"New best candidate: {candidate.name} with score {total_score:.2f}"
                    )

                    # Check for early stop
                    if total_score >= self.early_stop_threshold:
                        logger.info(
                            f"EARLY STOP: Found {candidate.name} with {total_score:.2f} confidence!"
                        )
                        self.found_answer.set()

                        if self.progress_callback:
                            self.progress_callback(
                                f"Found answer: {candidate.name} ({int(total_score * 100)}% confidence)",
                                95,
                                {
                                    "phase": "early_stop",
                                    "final_answer": candidate.name,
                                    "confidence": total_score,
                                },
                            )

            return total_score

        except Exception:
            logger.exception(f"Error evaluating candidate {candidate.name}")
            return 0.0

    def _progressive_constraint_search(self):
        """Override to implement early stopping."""
        current_candidates = []
        search_iterations = 0
        max_search_iterations = 3

        # Detect entity type
        self.entity_type = self._detect_entity_type()
        logger.info(f"Detected entity type: {self.entity_type}")

        while (
            search_iterations < max_search_iterations
            and not self.found_answer.is_set()
        ):
            search_iterations += 1

            # Create search combinations based on iteration
            if search_iterations == 1:
                combinations = self._create_strict_combinations()
                strictness = "strict"
            elif search_iterations == 2:
                combinations = self._create_relaxed_combinations()
                strictness = "relaxed"
            else:
                combinations = self._create_individual_combinations()
                strictness = "individual"

            logger.info(
                f"Iteration {search_iterations}: {strictness} mode with {len(combinations)} combinations"
            )

            # Run searches in parallel with immediate evaluation
            new_candidates = self._parallel_search(combinations)
            current_candidates.extend(new_candidates)

            # Check if we have enough results or found the answer
            unique_candidates = self._deduplicate_candidates(current_candidates)

            if self.found_answer.is_set():
                logger.info(f"Early stop - found answer: {self.best_candidate}")
                break

            if len(unique_candidates) >= self.min_results_threshold:
                logger.info(
                    f"Found {len(unique_candidates)} candidates - checking if we need more"
                )
                # Continue only if best score is below threshold
                if self.best_score >= 0.9:
                    logger.info(
                        f"Best score {self.best_score:.2f} is high enough - stopping search"
                    )
                    break

        # Set final candidates
        self.candidates = [
            c for c in unique_candidates if c.name == self.best_candidate
        ]
        if not self.candidates and unique_candidates:
            # If best candidate wasn't in the list somehow, use top scored candidates
            scored_candidates = sorted(
                unique_candidates,
                key=lambda c: self.evaluated_candidates.get(c.name, 0),
                reverse=True,
            )
            self.candidates = scored_candidates[: self.candidate_limit]

        self.final_answer = self.best_candidate
        self.confidence = self.best_score

    def analyze_topic(self, topic: str) -> Dict:
        """Analyze topic with early stopping."""
        # Call parent's analyze_topic to handle constraint extraction
        result = super().analyze_topic(topic)

        # Add our early stopping information
        result["early_stopped"] = self.found_answer.is_set()
        result["evaluated_candidates"] = self.evaluated_candidates
        result["best_candidate"] = self.best_candidate
        result["best_score"] = self.best_score

        return result

    def _gather_evidence_for_constraint(
        self, candidate: Candidate, constraint: Constraint
    ) -> List:
        """Gather evidence for a specific candidate-constraint pair."""
        # Run targeted search for this specific combination
        query = f'"{candidate.name}" {constraint.value} verification'

        try:
            results = self._execute_search(query)
            return self._extract_evidence_from_results(
                results, candidate, constraint
            )
        except Exception:
            logger.exception(f"Error gathering evidence for {candidate.name}")
            return []

    def _extract_evidence_from_results(
        self, results: Dict, candidate: Candidate, constraint: Constraint
    ) -> List:
        """Extract relevant evidence from search results."""
        evidence = []
        content = results.get("current_knowledge", "")

        if content:
            # Use LLM to extract evidence
            prompt = f"""
            Extract evidence regarding whether "{candidate.name}" satisfies this constraint:

            Constraint: {constraint.value}
            Constraint Type: {constraint.type.value}

            Search Results:
            {content[:3000]}

            Extract specific evidence that either supports or refutes the constraint.
            Return a confidence score from 0 to 1.
            """

            try:
                response = self.model.invoke(prompt).content
                evidence.append(
                    {
                        "text": response,
                        "source": "search_results",
                        "confidence": self._extract_confidence_from_response(
                            response
                        ),
                    }
                )
            except Exception:
                logger.exception("Error extracting evidence")

        return evidence

    def _extract_confidence_from_response(self, response: str) -> float:
        """Extract confidence score from LLM response."""
        # Simple extraction - look for number between 0 and 1
        import re

        pattern = r"\b0?\.\d+\b|\b1\.0\b|\b1\b"
        matches = re.findall(pattern, response)

        if matches:
            try:
                return float(matches[-1])  # Use last number found
            except ValueError:
                pass

        # Default confidence based on keywords
        if any(
            word in response.lower()
            for word in ["definitely", "certainly", "absolutely"]
        ):
            return 0.9
        if any(
            word in response.lower()
            for word in ["likely", "probably", "appears"]
        ):
            return 0.7
        if any(
            word in response.lower() for word in ["possibly", "maybe", "might"]
        ):
            return 0.5
        if any(
            word in response.lower() for word in ["unlikely", "doubtful", "not"]
        ):
            return 0.3

        return 0.5

    def _evaluate_evidence(
        self, evidence: List, constraint: Constraint
    ) -> float:
        """Evaluate evidence to determine constraint satisfaction score."""
        if not evidence:
            return 0.0

        # Average confidence across all evidence
        confidences = [e.get("confidence", 0.5) for e in evidence]
        return sum(confidences) / len(confidences) if confidences else 0.0
