"""
Improved evidence-based search strategy for complex query resolution.

Key improvements:
1. Multi-stage candidate discovery with adaptive query generation
2. Dynamic constraint combination for cross-constraint searches
3. Query adaptation based on partial results
4. Enhanced source diversity management
"""

import itertools
from collections import defaultdict
from dataclasses import dataclass
from datetime import datetime, UTC
from typing import Any, Dict, List, Set

from langchain_core.language_models import BaseChatModel

from ...utilities.search_utilities import remove_think_tags
from ..candidates.base_candidate import Candidate
from ..constraints.base_constraint import Constraint, ConstraintType
from ..constraints.constraint_analyzer import ConstraintAnalyzer
from ..evidence.base_evidence import EvidenceType
from ..evidence.evaluator import EvidenceEvaluator
from ..findings.repository import FindingsRepository
from .base_strategy import BaseSearchStrategy


@dataclass
class SearchAttempt:
    """Track search attempts for query adaptation."""

    query: str
    constraint_ids: List[str]
    results_count: int
    candidates_found: int
    timestamp: str
    strategy_type: str  # 'single', 'combined', 'exploratory'


class ImprovedEvidenceBasedStrategy(BaseSearchStrategy):
    """
    Improved evidence-based strategy with adaptive search capabilities.

    Key improvements:
    1. Multi-stage candidate discovery
    2. Adaptive query generation based on results
    3. Cross-constraint search optimization
    4. Source diversity tracking and enhancement
    """

    def __init__(
        self,
        model: BaseChatModel,
        search: Any,
        all_links_of_system: List[str],
        max_iterations: int = 20,
        confidence_threshold: float = 0.85,
        candidate_limit: int = 15,  # Increased for better diversity
        evidence_threshold: float = 0.6,
        max_search_iterations: int = 3,
        questions_per_iteration: int = 3,
        min_source_diversity: int = 3,  # Minimum different sources
        adaptive_query_count: int = 3,  # Number of adaptive queries per stage
    ):
        """Initialize the improved evidence-based strategy."""
        super().__init__(all_links_of_system)
        self.model = model
        self.search = search
        self.max_iterations = max_iterations
        self.confidence_threshold = confidence_threshold
        self.candidate_limit = candidate_limit
        self.evidence_threshold = evidence_threshold
        self.max_search_iterations = max_search_iterations
        self.questions_per_iteration = questions_per_iteration
        self.min_source_diversity = min_source_diversity
        self.adaptive_query_count = adaptive_query_count

        # Initialize components
        self.constraint_analyzer = ConstraintAnalyzer(model)
        self.evidence_evaluator = EvidenceEvaluator(model)
        self.findings_repository = FindingsRepository(model)

        # State tracking
        self.constraints: List[Constraint] = []
        self.candidates: List[Candidate] = []
        self.search_history: List[Dict] = []
        self.search_attempts: List[SearchAttempt] = []
        self.failed_queries: Set[str] = set()
        self.successful_patterns: List[Dict[str, Any]] = []
        self.source_types: Dict[str, Set[str]] = defaultdict(set)
        self.iteration: int = 0

    def analyze_topic(self, query: str) -> Dict:
        """Analyze a topic using improved evidence-based approach."""
        # Initialize
        self.all_links_of_system.clear()
        self.questions_by_iteration = []
        self.findings = []
        self.iteration = 0
        self.search_attempts.clear()
        self.failed_queries.clear()
        self.successful_patterns.clear()

        # Step 1: Extract initial constraints
        if self.progress_callback:
            self.progress_callback(
                "Analyzing query for constraint extraction...",
                2,
                {"phase": "constraint_analysis", "status": "starting"},
            )

        self.constraints = self.constraint_analyzer.extract_constraints(query)

        # Step 2: Multi-stage candidate discovery
        self._multi_stage_candidate_discovery()

        # Step 3: Main evidence-gathering loop with adaptive search
        while (
            self.iteration < self.max_iterations
            and not self._has_sufficient_answer()
        ):
            self.iteration += 1

            if self.progress_callback:
                progress = 15 + int((self.iteration / self.max_iterations) * 70)
                self.progress_callback(
                    f"Iteration {self.iteration}/{self.max_iterations} - {self._get_iteration_status()}",
                    progress,
                    {
                        "phase": "iteration_start",
                        "iteration": self.iteration,
                        "candidates_count": len(self.candidates),
                        "search_attempts": len(self.search_attempts),
                        "successful_patterns": len(self.successful_patterns),
                    },
                )

            # Adaptive evidence gathering
            self._adaptive_evidence_gathering()

            # Score and prune with diversity consideration
            self._score_with_diversity()

            # Adaptive candidate discovery if needed
            if len(self.candidates) < 3 or self._needs_diversity():
                self._adaptive_candidate_search()

        # Step 4: Cross-validation and final verification
        self._cross_validate_candidates()

        # Step 5: Generate final answer
        return self._synthesize_final_answer(query)

    def _multi_stage_candidate_discovery(self):
        """Multi-stage candidate discovery with different strategies."""
        stages = [
            ("distinctive", self._discover_with_distinctive_constraints),
            ("combined", self._discover_with_combined_constraints),
            ("exploratory", self._discover_with_exploratory_search),
            ("pattern_based", self._discover_with_pattern_matching),
        ]

        for stage_name, discovery_method in stages:
            if self.progress_callback:
                self.progress_callback(
                    f"Stage {stage_name}: Discovering candidates...",
                    5 + stages.index((stage_name, discovery_method)) * 3,
                    {"phase": "candidate_discovery", "stage": stage_name},
                )

            new_candidates = discovery_method()

            # Add unique candidates
            existing_names = {c.name.lower() for c in self.candidates}
            for candidate in new_candidates:
                if candidate.name.lower() not in existing_names:
                    self.candidates.append(candidate)
                    existing_names.add(candidate.name.lower())

            # Stop if we have enough diverse candidates
            if len(self.candidates) >= self.candidate_limit // 2:
                break

    def _discover_with_distinctive_constraints(self) -> List[Candidate]:
        """Discover candidates using most distinctive constraints."""
        distinctive = self._get_adaptive_distinctive_constraints()
        candidates = []

        for constraint_combo in self._generate_constraint_combinations(
            distinctive, max_size=3
        ):
            query = self._create_adaptive_search_query(constraint_combo)
            if query not in self.failed_queries:
                results = self._execute_tracked_search(
                    query, constraint_combo, "distinctive"
                )
                candidates.extend(
                    self._extract_candidates_with_context(results, query)
                )

                if candidates:  # Track successful patterns
                    self.successful_patterns.append(
                        {
                            "constraints": [c.id for c in constraint_combo],
                            "query_pattern": query,
                            "candidates_found": len(candidates),
                        }
                    )

        return candidates

    def _discover_with_combined_constraints(self) -> List[Candidate]:
        """Discover candidates using strategic constraint combinations."""
        # Combine constraints from different types for better results
        type_groups = defaultdict(list)
        for c in self.constraints:
            type_groups[c.type].append(c)

        candidates = []
        # Cross-type combinations
        for type1, type2 in itertools.combinations(type_groups.keys(), 2):
            for c1, c2 in itertools.product(
                type_groups[type1][:2], type_groups[type2][:2]
            ):
                query = self._create_cross_constraint_query([c1, c2])
                results = self._execute_tracked_search(
                    query, [c1, c2], "combined"
                )
                candidates.extend(
                    self._extract_candidates_with_context(results, query)
                )

        return candidates

    def _discover_with_exploratory_search(self) -> List[Candidate]:
        """Use exploratory searches to find unexpected candidates."""
        candidates = []

        # Generate exploratory queries
        exploratory_prompt = f"""
Based on these constraints, generate 3 exploratory search queries that might find relevant candidates:

Constraints:
{self._format_constraints_for_prompt(self.constraints[:5])}

Create queries that:
1. Use alternative phrasings or related concepts
2. Explore edge cases or unusual combinations
3. Look for historical or contextual matches

Return only the queries, one per line.
"""

        response = self.model.invoke(exploratory_prompt)
        queries = remove_think_tags(response.content).strip().split("\n")

        for query in queries[:3]:
            if query.strip() and query not in self.failed_queries:
                results = self._execute_tracked_search(
                    query, self.constraints[:3], "exploratory"
                )
                candidates.extend(
                    self._extract_candidates_with_context(results, query)
                )

        return candidates

    def _discover_with_pattern_matching(self) -> List[Candidate]:
        """Use pattern matching based on successful patterns."""
        if not self.successful_patterns:
            return []

        candidates = []

        # Adapt successful patterns
        for pattern in self.successful_patterns[:3]:
            constraint_ids = pattern["constraints"]
            constraints = [
                c for c in self.constraints if c.id in constraint_ids
            ]

            # Create variations of successful queries
            adapted_query = self._adapt_successful_query(
                pattern["query_pattern"], constraints
            )
            results = self._execute_tracked_search(
                adapted_query, constraints, "pattern_based"
            )
            candidates.extend(
                self._extract_candidates_with_context(results, adapted_query)
            )

        return candidates

    def _adaptive_evidence_gathering(self):
        """Gather evidence with adaptive query generation."""
        for candidate in self.candidates[:5]:
            unverified = candidate.get_unverified_constraints(self.constraints)

            if not unverified:
                continue

            # Sort by weight and group by type
            unverified.sort(key=lambda c: c.weight, reverse=True)
            type_groups = defaultdict(list)
            for c in unverified:
                type_groups[c.type].append(c)

            # Try different evidence gathering strategies
            for constraint_type, constraints in type_groups.items():
                # Try single constraint
                for c in constraints[:2]:
                    query = self._create_evidence_query(candidate, [c])
                    results = self._execute_tracked_search(
                        query, [c], "evidence"
                    )
                    evidence = self.evidence_evaluator.extract_evidence(
                        results.get("current_knowledge", ""), candidate.name, c
                    )
                    candidate.add_evidence(c.id, evidence)

                # Try combined constraints of same type
                if len(constraints) > 1:
                    query = self._create_evidence_query(
                        candidate, constraints[:2]
                    )
                    results = self._execute_tracked_search(
                        query, constraints[:2], "evidence_combined"
                    )

                    for c in constraints[:2]:
                        evidence = self.evidence_evaluator.extract_evidence(
                            results.get("current_knowledge", ""),
                            candidate.name,
                            c,
                        )
                        if (
                            c.id not in candidate.evidence
                            or evidence.confidence
                            > candidate.evidence[c.id].confidence
                        ):
                            candidate.add_evidence(c.id, evidence)

    def _create_adaptive_search_query(
        self, constraints: List[Constraint]
    ) -> str:
        """Create adaptive search queries based on past performance."""
        # Check if similar constraint combinations have been successful
        constraint_ids = {c.id for c in constraints}

        for pattern in self.successful_patterns:
            if (
                len(constraint_ids.intersection(pattern["constraints"]))
                >= len(constraint_ids) // 2
            ):
                # Adapt successful pattern
                return self._adapt_successful_query(
                    pattern["query_pattern"], constraints
                )

        # Check failed queries to avoid repetition
        base_query = self._create_base_search_query(constraints)
        if base_query in self.failed_queries:
            # Modify query to avoid failure
            return self._modify_failed_query(base_query, constraints)

        return base_query

    def _create_cross_constraint_query(
        self, constraints: List[Constraint]
    ) -> str:
        """Create queries that leverage relationships between constraints."""
        prompt = f"""
Create a search query that finds candidates satisfying BOTH/ALL of these constraints:

{self._format_constraints_for_prompt(constraints)}

The query should:
1. Find entities that match both/all constraints
2. Use operators to require both/all conditions
3. Focus on finding specific names or entities

Return only the search query.
"""

        response = self.model.invoke(prompt)
        return remove_think_tags(response.content).strip()

    def _create_evidence_query(
        self, candidate: Candidate, constraints: List[Constraint]
    ) -> str:
        """Create targeted evidence queries."""
        constraint_desc = self._format_constraints_for_prompt(constraints)

        prompt = f"""
Create a search query to verify if "{candidate.name}" satisfies these constraints:

{constraint_desc}

The query should:
1. Include the candidate name
2. Target the specific constraint requirements
3. Find factual evidence, not opinions

Return only the search query.
"""

        response = self.model.invoke(prompt)
        return remove_think_tags(response.content).strip()

    def _score_with_diversity(self):
        """Score candidates considering source diversity."""
        for candidate in self.candidates:
            # Base score from evidence
            candidate.calculate_score(self.constraints)

            # Diversity bonus
            diversity_score = self._calculate_diversity_score(candidate)
            candidate.score = 0.8 * candidate.score + 0.2 * diversity_score

            # Track source types
            for evidence in candidate.evidence.values():
                if hasattr(evidence, "source"):
                    self.source_types[candidate.name].add(evidence.source)

        # Sort by adjusted score
        self.candidates.sort(key=lambda c: c.score, reverse=True)

        # Prune while maintaining some diversity
        self._prune_with_diversity()

    def _cross_validate_candidates(self):
        """Cross-validate top candidates using different approaches."""
        if not self.candidates:
            return

        top_candidates = self.candidates[:3]

        for candidate in top_candidates:
            # Validate using different search engines or approaches
            validation_queries = self._generate_validation_queries(candidate)

            for query in validation_queries:
                results = self._execute_tracked_search(
                    query, self.constraints, "validation"
                )

                # Update evidence if better found
                for constraint in self.constraints:
                    evidence = self.evidence_evaluator.extract_evidence(
                        results.get("current_knowledge", ""),
                        candidate.name,
                        constraint,
                    )

                    if (
                        constraint.id not in candidate.evidence
                        or evidence.confidence
                        > candidate.evidence[constraint.id].confidence
                    ):
                        candidate.add_evidence(constraint.id, evidence)

    def _execute_tracked_search(
        self, query: str, constraints: List[Constraint], strategy_type: str
    ) -> Dict:
        """Execute search with tracking for adaptation."""
        results = self._execute_search(query)

        # Track the attempt
        candidates_found = len(
            self._extract_candidates_with_context(results, query)
        )
        attempt = SearchAttempt(
            query=query,
            constraint_ids=[c.id for c in constraints],
            results_count=len(results.get("all_links_of_system", [])),
            candidates_found=candidates_found,
            timestamp=datetime.now(UTC).isoformat(),
            strategy_type=strategy_type,
        )
        self.search_attempts.append(attempt)

        # Mark as failed if no results
        if candidates_found == 0:
            self.failed_queries.add(query)

        return results

    def _needs_diversity(self) -> bool:
        """Check if we need more diverse candidates."""
        if len(self.candidates) < 3:
            return True

        # Check source diversity
        top_sources = self.source_types.get(self.candidates[0].name, set())
        return len(top_sources) < self.min_source_diversity

    def _generate_constraint_combinations(
        self, constraints: List[Constraint], max_size: int = 3
    ) -> List[List[Constraint]]:
        """Generate strategic constraint combinations."""
        combinations = []

        # Single constraints
        combinations.extend([[c] for c in constraints])

        # Pairs
        for size in range(2, min(len(constraints), max_size) + 1):
            for combo in itertools.combinations(constraints, size):
                combinations.append(list(combo))

        return combinations

    def _format_constraints_for_prompt(
        self, constraints: List[Constraint]
    ) -> str:
        """Format constraints for LLM prompts."""
        formatted = []
        for c in constraints:
            formatted.append(
                f"- {c.type.value}: {c.description} (weight: {c.weight:.2f})"
            )
        return "\n".join(formatted)

    def _adapt_successful_query(
        self, pattern_query: str, constraints: List[Constraint]
    ) -> str:
        """Adapt a successful query pattern with new constraints."""
        prompt = f"""
Adapt this successful search query pattern with new constraints:

Original query: {pattern_query}

New constraints:
{self._format_constraints_for_prompt(constraints)}

Create a similar query structure but with the new constraint values.
Return only the adapted query.
"""

        response = self.model.invoke(prompt)
        return remove_think_tags(response.content).strip()

    def _modify_failed_query(
        self, failed_query: str, constraints: List[Constraint]
    ) -> str:
        """Modify a failed query to try a different approach."""
        prompt = f"""
This search query returned no results: {failed_query}

Constraints we're trying to satisfy:
{self._format_constraints_for_prompt(constraints)}

Create an alternative query that:
1. Uses different keywords or phrases
2. Tries a different search approach
3. Still targets the same constraints

Return only the modified query.
"""

        response = self.model.invoke(prompt)
        return remove_think_tags(response.content).strip()

    def _calculate_diversity_score(self, candidate: Candidate) -> float:
        """Calculate diversity score for a candidate."""
        if not candidate.evidence:
            return 0.0

        # Source diversity
        sources = self.source_types.get(candidate.name, set())
        source_score = min(len(sources) / self.min_source_diversity, 1.0)

        # Evidence type diversity
        evidence_types = {e.type for e in candidate.evidence.values()}
        type_score = len(evidence_types) / len(EvidenceType)

        # Confidence distribution (prefer balanced confidence)
        confidences = [e.confidence for e in candidate.evidence.values()]
        if confidences:
            variance = sum((c - 0.7) ** 2 for c in confidences) / len(
                confidences
            )
            confidence_score = 1.0 / (1.0 + variance)
        else:
            confidence_score = 0.0

        return (source_score + type_score + confidence_score) / 3.0

    def _prune_with_diversity(self):
        """Prune candidates while maintaining diversity."""
        if len(self.candidates) <= self.candidate_limit:
            return

        # Keep top candidates
        kept = self.candidates[: self.candidate_limit // 2]
        remaining = self.candidates[self.candidate_limit // 2 :]

        # Add diverse candidates from remaining
        for candidate in remaining:
            if len(kept) >= self.candidate_limit:
                break

            # Check if this candidate adds diversity
            if self._adds_diversity(candidate, kept):
                kept.append(candidate)

        # Fill remaining slots with highest scoring
        for candidate in remaining:
            if len(kept) >= self.candidate_limit:
                break
            if candidate not in kept:
                kept.append(candidate)

        self.candidates = kept

    def _adds_diversity(
        self, candidate: Candidate, existing: List[Candidate]
    ) -> bool:
        """Check if a candidate adds diversity to the existing set."""
        # Check source diversity
        candidate_sources = self.source_types.get(candidate.name, set())
        existing_sources = set()
        for c in existing:
            existing_sources.update(self.source_types.get(c.name, set()))

        new_sources = candidate_sources - existing_sources
        if new_sources:
            return True

        # Check constraint coverage
        candidate_constraints = set(candidate.evidence.keys())
        existing_constraints = set()
        for c in existing:
            existing_constraints.update(c.evidence.keys())

        new_constraints = candidate_constraints - existing_constraints
        return len(new_constraints) > 0

    def _generate_validation_queries(self, candidate: Candidate) -> List[str]:
        """Generate validation queries for cross-checking."""
        queries = []

        # Query combining multiple constraints
        high_weight_constraints = sorted(
            self.constraints, key=lambda c: c.weight, reverse=True
        )[:3]
        combined_query = f'"{candidate.name}" ' + " ".join(
            c.to_search_terms() for c in high_weight_constraints
        )
        queries.append(combined_query)

        # Query with alternative phrasing
        alt_prompt = f"""
Create an alternative search query to validate "{candidate.name}" as the answer.
Use different keywords but same intent.

Return only the query.
"""
        response = self.model.invoke(alt_prompt)
        queries.append(remove_think_tags(response.content).strip())

        # Source-specific query
        if self.source_types.get(candidate.name):
            source_query = f'"{candidate.name}" site:{list(self.source_types[candidate.name])[0]}'
            queries.append(source_query)

        return queries

    def _extract_candidates_with_context(
        self, results: Dict, query: str
    ) -> List[Candidate]:
        """Extract candidates with context awareness."""
        # Use the original extraction method but with context
        candidates = self._extract_candidates_from_results(results, query)

        # Add context metadata
        for candidate in candidates:
            candidate.metadata["discovery_query"] = query
            candidate.metadata["discovery_stage"] = self.iteration

        return candidates

    def _create_base_search_query(self, constraints: List[Constraint]) -> str:
        """Create a base search query from constraints."""
        # Use an improved prompt that considers constraint relationships
        prompt = f"""
Create a search query that finds specific entities satisfying these constraints:

{self._format_constraints_for_prompt(constraints)}

Guidelines:
1. Focus on finding names/entities, not general information
2. Use the most distinctive constraints
3. Combine constraints effectively
4. Keep the query concise but comprehensive

Return only the search query.
"""

        response = self.model.invoke(prompt)
        return remove_think_tags(response.content).strip()

    def _adaptive_candidate_search(self):
        """Adaptively search for more candidates based on current state."""
        # Analyze what types of candidates we're missing
        covered_constraints = set()
        for candidate in self.candidates[:5]:
            covered_constraints.update(candidate.evidence.keys())

        uncovered = [
            c for c in self.constraints if c.id not in covered_constraints
        ]

        if uncovered:
            # Search specifically for uncovered constraints
            queries = []
            for constraint_group in self._generate_constraint_combinations(
                uncovered[:5], max_size=2
            ):
                query = self._create_adaptive_search_query(constraint_group)
                queries.append((query, constraint_group))

            for query, constraints in queries[: self.adaptive_query_count]:
                results = self._execute_tracked_search(
                    query, constraints, "adaptive"
                )
                new_candidates = self._extract_candidates_with_context(
                    results, query
                )

                # Add unique candidates
                existing_names = {c.name.lower() for c in self.candidates}
                for candidate in new_candidates:
                    if candidate.name.lower() not in existing_names:
                        self.candidates.append(candidate)
                        existing_names.add(candidate.name.lower())

    def _get_adaptive_distinctive_constraints(self) -> List[Constraint]:
        """Get distinctive constraints with adaptive prioritization."""
        # Start with basic prioritization
        priority_order = [
            ConstraintType.NAME_PATTERN,
            ConstraintType.LOCATION,
            ConstraintType.EVENT,
            ConstraintType.STATISTIC,
            ConstraintType.COMPARISON,
            ConstraintType.PROPERTY,
            ConstraintType.TEMPORAL,
            ConstraintType.EXISTENCE,
        ]

        # Adjust based on successful patterns
        if self.successful_patterns:
            # Count successful constraint types
            type_success = defaultdict(int)
            for pattern in self.successful_patterns:
                for constraint_id in pattern["constraints"]:
                    constraint = next(
                        (c for c in self.constraints if c.id == constraint_id),
                        None,
                    )
                    if constraint:
                        type_success[constraint.type] += pattern[
                            "candidates_found"
                        ]

            # Sort by success rate
            priority_order = sorted(
                priority_order,
                key=lambda t: type_success.get(t, 0),
                reverse=True,
            )

        # Sort constraints by adjusted priority
        sorted_constraints = sorted(
            self.constraints,
            key=lambda c: (priority_order.index(c.type), -c.weight),
        )

        return sorted_constraints[:5]
