"""
Clustering Module - Union-Find algorithm for memory grouping.

Contains:
- UnionFind: Data structure for efficient set operations
- find_clusters: Discover clusters of similar memories
- tier_clusters: Categorize clusters by complexity
"""

from typing import Dict, List, Set, Tuple

from .models import Cluster, generate_cluster_id


class UnionFind:
    """Union-Find data structure for clustering."""

    def __init__(self):
        self.parent: Dict[str, str] = {}
        self.rank: Dict[str, int] = {}

    def find(self, x: str) -> str:
        """Find root with path compression."""
        if x not in self.parent:
            self.parent[x] = x
            self.rank[x] = 0
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: str, y: str) -> None:
        """Union by rank."""
        px, py = self.find(x), self.find(y)
        if px == py:
            return
        if self.rank[px] < self.rank[py]:
            px, py = py, px
        self.parent[py] = px
        if self.rank[px] == self.rank[py]:
            self.rank[px] += 1

    def get_groups(self) -> Dict[str, Set[str]]:
        """Get all groups."""
        groups: Dict[str, Set[str]] = {}
        for x in self.parent:
            root = self.find(x)
            if root not in groups:
                groups[root] = set()
            groups[root].add(x)
        return groups


def find_clusters(
    store,
    distance_threshold: float = 0.5
) -> List[Cluster]:
    """
    Find clusters of similar memories using union-find.

    Algorithm:
    1. Get all memories with their vectors
    2. For each memory, find all others within threshold distance
    3. Build adjacency graph using union-find
    4. Extract connected components (clusters)
    5. Return clusters with 2+ members

    Args:
        store: ContextStore instance
        distance_threshold: Max L2 distance to consider similar (default 0.55)

    Returns:
        List of Cluster objects with 2+ members
    """
    # Get all memories with vectors
    memories = store.get_all_memories_with_vectors()
    if len(memories) < 2:
        return []

    # Build memory lookup
    memory_by_id = {m["id"]: m for m in memories}

    # Union-find for clustering
    uf = UnionFind()

    # Track distances between pairs for average calculation
    pair_distances: Dict[Tuple[str, str], float] = {}

    # For each memory, find similar ones
    print(f"[Consolidate] Scanning {len(memories)} memories for clusters...")

    for mem in memories:
        mem_id = mem["id"]

        # Find similar memories using the store's search
        similar = store.find_similar_memories(
            memory_id=mem_id,
            threshold=distance_threshold,
            limit=50
        )

        # Union with all similar memories
        for sim in similar:
            sim_id = sim["id"]
            uf.union(mem_id, sim_id)

            # Record distance for average calculation
            pair_key = tuple(sorted([mem_id, sim_id]))
            if pair_key not in pair_distances:
                pair_distances[pair_key] = sim["distance"]

    # Extract groups with 2+ members
    groups = uf.get_groups()
    clusters = []

    for root, member_ids in groups.items():
        if len(member_ids) >= 2:
            # Get full memory data for each member
            cluster_memories = [memory_by_id[mid] for mid in member_ids if mid in memory_by_id]

            # Calculate average distance within cluster
            cluster_pairs = [(a, b) for a in member_ids for b in member_ids if a < b]
            distances = [pair_distances.get(tuple(sorted([a, b])), 0) for a, b in cluster_pairs]
            avg_dist = sum(distances) / len(distances) if distances else 0.0

            # Generate stable hash-based cluster ID immediately
            cluster_id = generate_cluster_id(list(member_ids))

            clusters.append(Cluster(
                id=cluster_id,
                memory_ids=list(member_ids),
                memories=cluster_memories,
                avg_distance=avg_dist
            ))

    # Sort by size (largest first), then by average distance (closest first),
    # then by ID for deterministic ordering when size and distance are equal
    clusters.sort(key=lambda c: (-c.size, c.avg_distance, c.id))

    print(f"[Consolidate] Found {len(clusters)} clusters")
    return clusters


def tier_clusters(clusters: List[Cluster]) -> Dict[str, List[Cluster]]:
    """
    Tier clusters by complexity for processing.

    All clusters now go to LLM - no auto tier.
    Length-based heuristics are unreliable: shorter memories can be more
    accurate/current, longer ones can be verbose but less useful.

    Tiers:
    - auto: Always empty (no auto-processing)
    - llm_small: Clusters of 2-5 entries
    - llm_large: Large clusters (5+) needing full analysis

    Returns:
        Dict with keys: "auto", "llm_small", "llm_large"
    """
    tiered = {
        "auto": [],       # Always empty - no auto processing
        "llm_small": [],  # Clusters of 2-5
        "llm_large": []   # Clusters of 5+
    }

    for cluster in clusters:
        if cluster.size <= 5:
            tiered["llm_small"].append(cluster)
        else:
            tiered["llm_large"].append(cluster)

    return tiered
