"""
Multi-Stage Constraint Verification with Intelligent Weighting

This module implements the constraint satisfaction tracking system outlined in
BROWSECOMP_IMPROVEMENT_STRATEGY.md to improve BrowseComp performance.

Key features:
1. Intelligent constraint weighting based on reliability
2. Progressive constraint verification
3. Partial satisfaction scoring
4. Constraint difficulty analysis
"""

from dataclasses import dataclass
from typing import Dict, List

from loguru import logger


@dataclass
class ConstraintVerificationResult:
    """Result of constraint verification process."""

    total_score: float
    should_accept: bool
    constraint_details: Dict[str, Dict]
    weighted_score: float
    reliability_adjusted_score: float


class ConstraintSatisfactionTracker:
    """
    Multi-stage constraint verification with intelligent weighting.

    Implements the strategy from BROWSECOMP_IMPROVEMENT_STRATEGY.md to:
    - Weight constraints by reliability and importance
    - Allow partial satisfaction instead of requiring perfect matches
    - Prioritize reliable constraints over problematic ones
    """

    def __init__(self, satisfaction_threshold: float = 0.7):
        self.satisfaction_threshold = satisfaction_threshold

        # Constraint type weights based on reliability and importance
        self.constraint_weights = {
            "PROPERTY": 1.0,  # Basic properties, moderately reliable
            "NAME_PATTERN": 1.5,  # Names are highly identifying
            "STATISTIC": 0.8,  # Often imprecise/outdated
            "EVENT": 1.2,  # Events are good identifiers
            "TEMPORAL": 1.1,  # Dates are useful but sometimes fuzzy
            "LOCATION": 1.3,  # Geographic info is reliable
            "COMPARISON": 0.6,  # Relative comparisons often problematic
            "EXISTENCE": 1.0,  # Basic existence checks
            "RELATIONSHIP": 0.9,  # Relationships can be complex
        }

        # Difficulty multipliers (easier constraints get higher weight)
        self.difficulty_multipliers = {
            "EASY": 1.2,  # Simple property checks
            "MEDIUM": 1.0,  # Numerical comparisons
            "HARD": 0.7,  # Complex relationships, ratios
        }

    def verify_constraints_progressively(
        self, candidate: object, constraints: List[object]
    ) -> ConstraintVerificationResult:
        """
        Score each constraint independently with intelligent weighting.

        Args:
            candidate: The candidate entity to verify
            constraints: List of constraints to check

        Returns:
            ConstraintVerificationResult with scoring details
        """
        scores = {}
        total_weight = 0

        for constraint in constraints:
            # Analyze constraint difficulty
            difficulty = self.analyze_constraint_difficulty(constraint)

            # Calculate constraint weight
            weight = self.calculate_constraint_weight(constraint, difficulty)

            # Verify constraint against candidate
            satisfaction = self.verify_single_constraint(candidate, constraint)

            # Store detailed results
            scores[
                constraint.id if hasattr(constraint, "id") else str(constraint)
            ] = {
                "satisfaction": satisfaction,
                "weight": weight,
                "difficulty": difficulty,
                "constraint_type": self._get_constraint_type(constraint),
                "raw_score": satisfaction,
                "weighted_score": satisfaction * weight,
            }

            total_weight += weight

        # Calculate weighted satisfaction score
        if total_weight > 0:
            weighted_score = (
                sum(
                    score["satisfaction"] * score["weight"]
                    for score in scores.values()
                )
                / total_weight
            )
        else:
            weighted_score = 0.0

        # Apply reliability adjustment
        reliability_adjusted_score = self._apply_reliability_adjustment(
            weighted_score, scores
        )

        # Determine if candidate should be accepted
        should_accept = (
            reliability_adjusted_score >= self.satisfaction_threshold
        )

        return ConstraintVerificationResult(
            total_score=reliability_adjusted_score,
            should_accept=should_accept,
            constraint_details=scores,
            weighted_score=weighted_score,
            reliability_adjusted_score=reliability_adjusted_score,
        )

    def analyze_constraint_difficulty(self, constraint: object) -> str:
        """
        Classify constraints by difficulty level.

        Returns:
            'EASY', 'MEDIUM', or 'HARD'
        """
        constraint_text = str(constraint).lower()

        # HARD: Complex relationships and ratios
        hard_indicators = [
            "times more",
            "times larger",
            "times bigger",
            "ratio",
            "proportion",
            "percentage of",
            "compared to",
            "relative to",
            "in relation to",
            "between",
            "among",
            "relationship",
        ]

        if any(indicator in constraint_text for indicator in hard_indicators):
            return "HARD"

        # MEDIUM: Numerical comparisons and ranges
        medium_indicators = [
            "number",
            "count",
            "amount",
            "quantity",
            "more than",
            "less than",
            "at least",
            "at most",
            "approximately",
            "around",
            "about",
            "million",
            "billion",
            "thousand",
        ]

        if any(indicator in constraint_text for indicator in medium_indicators):
            return "MEDIUM"

        # EASY: Simple property checks
        return "EASY"

    def calculate_constraint_weight(
        self, constraint: object, difficulty: str
    ) -> float:
        """
        Calculate the weight for a constraint based on type and difficulty.

        Args:
            constraint: The constraint object
            difficulty: Difficulty level ('EASY', 'MEDIUM', 'HARD')

        Returns:
            Float weight value
        """
        # Get base weight from constraint type
        constraint_type = self._get_constraint_type(constraint)
        base_weight = self.constraint_weights.get(constraint_type, 1.0)

        # Apply difficulty multiplier
        difficulty_multiplier = self.difficulty_multipliers.get(difficulty, 1.0)

        # Calculate final weight
        final_weight = base_weight * difficulty_multiplier

        logger.debug(
            f"Constraint weight: {constraint_type} × {difficulty} = "
            f"{base_weight} × {difficulty_multiplier} = {final_weight}"
        )

        return final_weight

    def verify_single_constraint(
        self, candidate: object, constraint: object
    ) -> float:
        """
        Verify a single constraint against a candidate.

        This is a placeholder - implement actual verification logic
        based on your constraint checking system.

        Returns:
            Float score from 0.0 to 1.0
        """
        # TODO: Implement actual constraint verification
        # This should call your existing constraint checker

        # For now, return a placeholder score
        # In practice, this would call something like:
        # return self.constraint_checker.verify(candidate, constraint)

        return 0.5  # Placeholder

    def _get_constraint_type(self, constraint: object) -> str:
        """Extract constraint type from constraint object."""
        if hasattr(constraint, "type"):
            if hasattr(constraint.type, "value"):
                return constraint.type.value
            else:
                return str(constraint.type)
        elif hasattr(constraint, "constraint_type"):
            return constraint.constraint_type
        else:
            # Try to infer from constraint text
            constraint_text = str(constraint).lower()

            if any(
                word in constraint_text
                for word in ["name", "called", "known as"]
            ):
                return "NAME_PATTERN"
            elif any(
                word in constraint_text
                for word in ["location", "country", "city"]
            ):
                return "LOCATION"
            elif any(
                word in constraint_text
                for word in ["year", "date", "when", "time"]
            ):
                return "TEMPORAL"
            elif any(
                word in constraint_text
                for word in ["number", "count", "amount"]
            ):
                return "STATISTIC"
            elif any(
                word in constraint_text
                for word in ["event", "happened", "occurred"]
            ):
                return "EVENT"
            elif any(
                word in constraint_text
                for word in ["than", "more", "less", "compared"]
            ):
                return "COMPARISON"
            else:
                return "PROPERTY"

    def _apply_reliability_adjustment(
        self, weighted_score: float, constraint_scores: Dict[str, Dict]
    ) -> float:
        """
        Apply reliability adjustment based on constraint mix.

        Boosts score if high-reliability constraints are satisfied,
        reduces score if only low-reliability constraints match.
        """
        # Count high vs low reliability constraints
        high_reliability_count = 0
        high_reliability_satisfied = 0
        low_reliability_count = 0
        low_reliability_satisfied = 0

        high_reliability_types = {
            "NAME_PATTERN",
            "LOCATION",
            "EVENT",
            "TEMPORAL",
        }
        low_reliability_types = {"COMPARISON", "STATISTIC"}

        for constraint_id, score_data in constraint_scores.items():
            constraint_type = score_data.get("constraint_type", "PROPERTY")
            satisfaction = score_data.get("satisfaction", 0.0)

            if constraint_type in high_reliability_types:
                high_reliability_count += 1
                if satisfaction > 0.5:
                    high_reliability_satisfied += 1
            elif constraint_type in low_reliability_types:
                low_reliability_count += 1
                if satisfaction > 0.5:
                    low_reliability_satisfied += 1

        # Calculate reliability bonus/penalty
        reliability_adjustment = 0.0

        # Bonus for satisfying high-reliability constraints
        if high_reliability_count > 0:
            high_reliability_ratio = (
                high_reliability_satisfied / high_reliability_count
            )
            reliability_adjustment += high_reliability_ratio * 0.1

        # Small penalty if only low-reliability constraints are satisfied
        if low_reliability_count > 0 and high_reliability_satisfied == 0:
            low_reliability_ratio = (
                low_reliability_satisfied / low_reliability_count
            )
            if low_reliability_ratio > 0.7:  # Many low-reliability matches
                reliability_adjustment -= 0.05

        # Apply adjustment
        adjusted_score = weighted_score + reliability_adjustment

        # Ensure score stays in valid range
        return max(0.0, min(1.0, adjusted_score))

    def get_constraint_priorities(self) -> Dict[str, int]:
        """
        Get constraint priorities for relaxation strategies.

        Returns priorities from 1 (relax first) to 10 (never relax).
        """
        return {
            "COMPARISON": 1,  # Frequently relax
            "STATISTIC": 3,  # Often relax
            "PROPERTY": 6,
            "EVENT": 5,
            "TEMPORAL": 7,
            "LOCATION": 8,
            "EXISTENCE": 9,  # Rarely relax
            "NAME_PATTERN": 10,  # Never relax
        }

    def suggest_constraint_relaxation(
        self, constraints: List[object], current_results: List[object]
    ) -> List[List[object]]:
        """
        Suggest progressive constraint relaxation based on results.

        Returns list of relaxed constraint sets to try.
        """
        if len(current_results) > 5:
            return [constraints]  # Enough results, no relaxation needed

        priorities = self.get_constraint_priorities()

        # Sort constraints by relaxation priority (lowest first)
        relaxable_constraints = sorted(
            constraints,
            key=lambda c: priorities.get(self._get_constraint_type(c), 5),
        )

        # Generate progressive relaxation sets
        relaxed_sets = []
        for i in range(1, min(len(constraints), 4)):  # Max 3 relaxation levels
            relaxed_set = relaxable_constraints[:-i]  # Remove i lowest priority
            if len(relaxed_set) >= 2:  # Keep at least 2 constraints
                relaxed_sets.append(relaxed_set)

        return relaxed_sets
