"""
Cross-constraint search optimization manager.
"""

import itertools
from collections import defaultdict
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Set, Tuple

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


@dataclass
class ConstraintRelationship:
    """Represents a relationship between constraints."""

    constraint1_id: str
    constraint2_id: str
    relationship_type: str  # 'complementary', 'dependent', 'exclusive'
    strength: float  # 0.0 to 1.0
    evidence: List[str] = field(default_factory=list)


@dataclass
class ConstraintCluster:
    """Group of related constraints that should be searched together."""

    constraints: List[Constraint]
    cluster_type: str  # 'temporal', 'spatial', 'causal', 'descriptive'
    coherence_score: float
    search_queries: List[str] = field(default_factory=list)


class CrossConstraintManager:
    """
    Manages cross-constraint relationships and optimizes multi-constraint searches.

    Key features:
    1. Identifies relationships between constraints
    2. Clusters related constraints for efficient searching
    3. Generates cross-constraint validation queries
    4. Tracks cross-constraint evidence patterns
    """

    def __init__(self, model: BaseChatModel):
        """Initialize the cross-constraint manager."""
        self.model = model
        self.relationships: Dict[Tuple[str, str], ConstraintRelationship] = {}
        self.clusters: List[ConstraintCluster] = []
        self.cross_validation_patterns: Dict[str, List[Dict]] = defaultdict(
            list
        )
        self.constraint_graph: Dict[str, Set[str]] = defaultdict(set)

    def analyze_constraint_relationships(
        self, constraints: List[Constraint]
    ) -> Dict[Tuple[str, str], ConstraintRelationship]:
        """Analyze relationships between constraints."""
        relationships = {}

        # Analyze each pair of constraints
        for c1, c2 in itertools.combinations(constraints, 2):
            relationship = self._analyze_pair(c1, c2)
            if (
                relationship.strength > 0.3
            ):  # Only keep meaningful relationships
                key = (c1.id, c2.id)
                relationships[key] = relationship

                # Update constraint graph
                self.constraint_graph[c1.id].add(c2.id)
                self.constraint_graph[c2.id].add(c1.id)

        self.relationships.update(relationships)
        return relationships

    def _analyze_pair(
        self, c1: Constraint, c2: Constraint
    ) -> ConstraintRelationship:
        """Analyze the relationship between two constraints."""
        prompt = f"""
Analyze the relationship between these two constraints:

Constraint 1: {c1.description} (Type: {c1.type.value})
Constraint 2: {c2.description} (Type: {c2.type.value})

Determine:
1. Relationship type (complementary, dependent, exclusive, or none)
2. Strength of relationship (0.0 to 1.0)
3. Brief explanation

Format:
Type: [relationship_type]
Strength: [0.0-1.0]
Evidence: [explanation]
"""

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

        # Parse response
        rel_type = "none"
        strength = 0.0
        evidence = []

        for line in content.strip().split("\n"):
            if line.startswith("Type:"):
                rel_type = line.split(":", 1)[1].strip().lower()
            elif line.startswith("Strength:"):
                try:
                    strength = float(line.split(":", 1)[1].strip())
                except ValueError:
                    strength = 0.0
            elif line.startswith("Evidence:"):
                evidence.append(line.split(":", 1)[1].strip())

        return ConstraintRelationship(
            constraint1_id=c1.id,
            constraint2_id=c2.id,
            relationship_type=rel_type,
            strength=strength,
            evidence=evidence,
        )

    def create_constraint_clusters(
        self, constraints: List[Constraint]
    ) -> List[ConstraintCluster]:
        """Create clusters of related constraints."""
        # First, analyze relationships if not done
        if not self.relationships:
            self.analyze_constraint_relationships(constraints)

        # Create clusters using different strategies
        clusters = []

        # 1. Type-based clusters
        type_groups = defaultdict(list)
        for c in constraints:
            type_groups[c.type].append(c)

        for group in type_groups.values():
            if len(group) > 1:
                cluster = ConstraintCluster(
                    constraints=group,
                    cluster_type="type_based",
                    coherence_score=0.7,
                )
                clusters.append(cluster)

        # 2. Relationship-based clusters
        strong_relationships = [
            rel for rel in self.relationships.values() if rel.strength > 0.6
        ]

        relationship_clusters = self._create_relationship_clusters(
            constraints, strong_relationships
        )
        clusters.extend(relationship_clusters)

        # 3. Semantic clusters
        semantic_clusters = self._create_semantic_clusters(constraints)
        clusters.extend(semantic_clusters)

        # Remove duplicate clusters
        unique_clusters = self._deduplicate_clusters(clusters)

        self.clusters = unique_clusters
        return unique_clusters

    def _create_relationship_clusters(
        self,
        constraints: List[Constraint],
        relationships: List[ConstraintRelationship],
    ) -> List[ConstraintCluster]:
        """Create clusters based on strong relationships."""
        clusters = []
        processed = set()

        # Build adjacency list
        adj_list = defaultdict(list)
        for rel in relationships:
            adj_list[rel.constraint1_id].append(rel.constraint2_id)
            adj_list[rel.constraint2_id].append(rel.constraint1_id)

        # Find connected components
        for constraint in constraints:
            if constraint.id in processed:
                continue

            # BFS to find connected component
            component = []
            queue = [constraint.id]
            visited = {constraint.id}

            while queue:
                current_id = queue.pop(0)
                current = next(
                    (c for c in constraints if c.id == current_id), None
                )
                if current:
                    component.append(current)
                    processed.add(current_id)

                    for neighbor_id in adj_list[current_id]:
                        if neighbor_id not in visited:
                            visited.add(neighbor_id)
                            queue.append(neighbor_id)

            if len(component) > 1:
                cluster = ConstraintCluster(
                    constraints=component,
                    cluster_type="relationship_based",
                    coherence_score=self._calculate_cluster_coherence(
                        component
                    ),
                )
                clusters.append(cluster)

        return clusters

    def _create_semantic_clusters(
        self, constraints: List[Constraint]
    ) -> List[ConstraintCluster]:
        """Create clusters based on semantic similarity."""
        prompt = f"""
Group these constraints into semantic clusters based on their meaning and intent:

{self._format_constraints_for_clustering(constraints)}

For each cluster:
1. List the constraint IDs
2. Describe the cluster theme
3. Rate coherence (0.0-1.0)

Format:
CLUSTER_1:
Constraints: [id1, id2, ...]
Theme: [description]
Coherence: [0.0-1.0]
"""

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

        clusters = []
        current_cluster = {}

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

            if line.startswith("CLUSTER_"):
                if current_cluster and "constraints" in current_cluster:
                    # Create cluster from previous data
                    constraint_ids = current_cluster["constraints"]
                    cluster_constraints = [
                        c for c in constraints if c.id in constraint_ids
                    ]

                    if len(cluster_constraints) > 1:
                        cluster = ConstraintCluster(
                            constraints=cluster_constraints,
                            cluster_type="semantic",
                            coherence_score=float(
                                current_cluster.get("coherence", 0.5)
                            ),
                        )
                        clusters.append(cluster)

                current_cluster = {}

            elif line.startswith("Constraints:"):
                ids_str = line.split(":", 1)[1].strip()
                # Extract IDs from various formats
                import re

                ids = re.findall(r"c\d+", ids_str)
                current_cluster["constraints"] = ids

            elif line.startswith("Theme:"):
                current_cluster["theme"] = line.split(":", 1)[1].strip()

            elif line.startswith("Coherence:"):
                try:
                    current_cluster["coherence"] = float(
                        line.split(":", 1)[1].strip()
                    )
                except ValueError:
                    current_cluster["coherence"] = 0.5

        # Don't forget the last cluster
        if current_cluster and "constraints" in current_cluster:
            constraint_ids = current_cluster["constraints"]
            cluster_constraints = [
                c for c in constraints if c.id in constraint_ids
            ]

            if len(cluster_constraints) > 1:
                cluster = ConstraintCluster(
                    constraints=cluster_constraints,
                    cluster_type="semantic",
                    coherence_score=float(
                        current_cluster.get("coherence", 0.5)
                    ),
                )
                clusters.append(cluster)

        return clusters

    def generate_cross_constraint_queries(
        self, cluster: ConstraintCluster
    ) -> List[str]:
        """Generate optimized queries for a constraint cluster."""
        queries = []

        # 1. Combined query (all constraints)
        combined_query = self._generate_combined_query(cluster.constraints)
        queries.append(combined_query)

        # 2. Progressive queries (build up constraints)
        progressive_queries = self._generate_progressive_queries(
            cluster.constraints
        )
        queries.extend(progressive_queries)

        # 3. Intersection queries (shared aspects)
        intersection_query = self._generate_intersection_query(
            cluster.constraints
        )
        if intersection_query:
            queries.append(intersection_query)

        # 4. Validation queries (cross-check)
        validation_queries = self._generate_validation_queries(
            cluster.constraints
        )
        queries.extend(validation_queries)

        # Store queries in cluster
        cluster.search_queries = queries

        return queries

    def _generate_combined_query(self, constraints: List[Constraint]) -> str:
        """Generate a query combining all constraints."""
        prompt = f"""
Create a search query that finds entities satisfying ALL of these related constraints:

{self._format_constraints_for_query(constraints)}

The query should:
1. Efficiently combine all constraints
2. Use appropriate operators (AND, OR)
3. Focus on finding specific entities
4. Be neither too broad nor too narrow

Return only the search query.
"""

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

    def _generate_progressive_queries(
        self, constraints: List[Constraint]
    ) -> List[str]:
        """Generate queries that progressively add constraints."""
        queries = []

        # Sort by weight/importance
        sorted_constraints = sorted(
            constraints, key=lambda c: c.weight, reverse=True
        )

        # Build up constraints
        for i in range(2, min(len(sorted_constraints) + 1, 4)):
            subset = sorted_constraints[:i]
            query = self._generate_combined_query(subset)
            queries.append(query)

        return queries

    def _generate_intersection_query(
        self, constraints: List[Constraint]
    ) -> Optional[str]:
        """Generate a query focused on the intersection of constraints."""
        if len(constraints) < 2:
            return None

        prompt = f"""
Identify the common theme or intersection among these constraints:

{self._format_constraints_for_query(constraints)}

Create a search query that targets this common aspect.
Return only the search query, or 'NONE' if no clear intersection exists.
"""

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

        if query.upper() == "NONE":
            return None

        return query

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

        # Pairwise validation queries
        for c1, c2 in itertools.combinations(constraints[:3], 2):
            prompt = f"""
Create a validation query that checks if an entity satisfies both:
- {c1.description}
- {c2.description}

Return only the search query.
"""

            response = self.model.invoke(prompt)
            query = remove_think_tags(response.content).strip()
            queries.append(query)

        return queries[:2]  # Limit to 2 validation queries

    def validate_candidate_across_constraints(
        self, candidate: Candidate, constraints: List[Constraint]
    ) -> Dict[str, float]:
        """Validate a candidate across multiple constraints simultaneously."""
        validation_scores = {}

        # Find relevant clusters for these constraints
        relevant_clusters = [
            cluster
            for cluster in self.clusters
            if any(c in cluster.constraints for c in constraints)
        ]

        for cluster in relevant_clusters:
            # Use cluster-specific queries for validation
            cluster_score = self._validate_with_cluster(candidate, cluster)

            # Update individual constraint scores
            for constraint in cluster.constraints:
                if constraint in constraints:
                    validation_scores[constraint.id] = max(
                        validation_scores.get(constraint.id, 0.0), cluster_score
                    )

        # Additional pairwise validation
        for c1, c2 in itertools.combinations(constraints, 2):
            if (c1.id, c2.id) in self.relationships:
                rel = self.relationships[(c1.id, c2.id)]
                if rel.relationship_type == "complementary":
                    # Boost scores for complementary constraints
                    pair_score = self._validate_pair(candidate, c1, c2)
                    validation_scores[c1.id] = max(
                        validation_scores.get(c1.id, 0.0), pair_score
                    )
                    validation_scores[c2.id] = max(
                        validation_scores.get(c2.id, 0.0), pair_score
                    )

        return validation_scores

    def _validate_with_cluster(
        self, candidate: Candidate, cluster: ConstraintCluster
    ) -> float:
        """Validate candidate using cluster-based approach."""
        if not cluster.search_queries:
            cluster.search_queries = self.generate_cross_constraint_queries(
                cluster
            )

        # Use the most comprehensive query
        validation_query = cluster.search_queries[0]

        prompt = f"""
Does "{candidate.name}" satisfy this multi-constraint query:
Query: {validation_query}

Constraints being checked:
{self._format_constraints_for_query(cluster.constraints)}

Provide a confidence score (0.0-1.0) based on how well the candidate matches.

Format:
Score: [0.0-1.0]
Explanation: [brief explanation]
"""

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

        # Parse score
        score = 0.0
        for line in content.strip().split("\n"):
            if line.startswith("Score:"):
                try:
                    score = float(line.split(":", 1)[1].strip())
                except ValueError:
                    score = 0.0
                break

        return score

    def _validate_pair(
        self, candidate: Candidate, c1: Constraint, c2: Constraint
    ) -> float:
        """Validate candidate against a pair of constraints."""
        prompt = f"""
Evaluate if "{candidate.name}" satisfies BOTH constraints:

1. {c1.description} (Type: {c1.type.value})
2. {c2.description} (Type: {c2.type.value})

Consider how these constraints relate to each other and whether the candidate satisfies both.

Provide a confidence score (0.0-1.0).

Format:
Score: [0.0-1.0]
"""

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

        # Parse score
        score = 0.0
        for line in content.strip().split("\n"):
            if line.startswith("Score:"):
                try:
                    score = float(line.split(":", 1)[1].strip())
                except ValueError:
                    score = 0.0
                break

        return score

    def _calculate_cluster_coherence(
        self, constraints: List[Constraint]
    ) -> float:
        """Calculate coherence score for a constraint cluster."""
        if len(constraints) < 2:
            return 0.0

        # Calculate based on relationship strengths
        total_strength = 0.0
        pair_count = 0

        for c1, c2 in itertools.combinations(constraints, 2):
            key = (c1.id, c2.id)
            if key in self.relationships:
                total_strength += self.relationships[key].strength
                pair_count += 1

        if pair_count == 0:
            return 0.5  # Default coherence

        average_strength = total_strength / pair_count

        # Adjust for cluster size (larger clusters with high average strength are better)
        size_factor = min(len(constraints) / 5.0, 1.0)

        return average_strength * (0.7 + 0.3 * size_factor)

    def _deduplicate_clusters(
        self, clusters: List[ConstraintCluster]
    ) -> List[ConstraintCluster]:
        """Remove duplicate clusters."""
        unique_clusters = []
        seen_sets = []

        for cluster in clusters:
            constraint_set = {c.id for c in cluster.constraints}

            # Check if we've seen this set
            is_duplicate = False
            for seen_set in seen_sets:
                if constraint_set == seen_set:
                    is_duplicate = True
                    break

            if not is_duplicate:
                unique_clusters.append(cluster)
                seen_sets.append(constraint_set)

        return unique_clusters

    def _format_constraints_for_clustering(
        self, constraints: List[Constraint]
    ) -> str:
        """Format constraints for clustering prompt."""
        formatted = []
        for c in constraints:
            formatted.append(
                f"{c.id}: {c.description} (Type: {c.type.value}, Weight: {c.weight})"
            )
        return "\n".join(formatted)

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

    def optimize_search_order(
        self, clusters: List[ConstraintCluster]
    ) -> List[ConstraintCluster]:
        """Optimize the order in which clusters should be searched."""
        # Sort by coherence and cluster size
        return sorted(
            clusters,
            key=lambda c: c.coherence_score * len(c.constraints),
            reverse=True,
        )
