"""
RBSP Semantic Resolver.

Uses LSH clusters for semantic similarity resolution.  Given a query
signal's semantic fingerprint, retrieves the corresponding LSH bucket and
expands to neighboring buckets for broader semantic coverage.

This module bridges the gap between exact hash matching and full semantic
similarity, providing O(1) amortised semantic resolution without AI models
or embedding vectors.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, Hashable, List, Optional, Set, Tuple

from rbsp.core.lsh import LSHIndex, MinHasher, ngram_shingles, word_shingles
from rbsp.query.signal import Signal


@dataclass(frozen=True, slots=True)
class SemanticResult:
    """A single semantically matched result.

    Attributes:
        item_id:    Unique identifier of the matched item.
        similarity: Estimated Jaccard similarity in [0.0, 1.0].
        bucket_id:  The LSH band/bucket where the match was found.
        metadata:   Optional metadata associated with the item.
    """
    item_id: Hashable
    similarity: float
    bucket_id: int = 0
    metadata: Optional[Dict[str, Any]] = None


class SemanticCluster:
    """A cluster of semantically related items sharing an LSH bucket.

    Attributes:
        cluster_id: Unique identifier for this cluster (band_idx, band_hash).
        members:    Set of item_ids in this cluster.
    """

    def __init__(self, cluster_id: Tuple[int, int]) -> None:
        self.cluster_id = cluster_id
        self.members: Set[Hashable] = set()

    def add(self, item_id: Hashable) -> None:
        """Add an item to this cluster."""
        self.members.add(item_id)

    def remove(self, item_id: Hashable) -> None:
        """Remove an item from this cluster."""
        self.members.discard(item_id)

    def __len__(self) -> int:
        return len(self.members)

    def __contains__(self, item_id: Hashable) -> bool:
        return item_id in self.members


class SemanticResolver:
    """Resolves queries to semantically related items via LSH.

    The resolver wraps an LSHIndex and provides a higher-level interface
    for semantic search, including:

        - Direct bucket lookup for the query's semantic fingerprint
        - Neighbor-bucket expansion for broader recall
        - Cluster-based result grouping
        - Configurable similarity thresholds

    Args:
        lsh_index:        The LSHIndex backing semantic resolution.
        expansion_depth:  Number of neighboring bands to check (default 1).
                          Higher values increase recall at the cost of speed.
        min_similarity:   Minimum Jaccard similarity to include in results.
        ngram_size:       Character n-gram size for query shingling.
    """

    def __init__(
        self,
        lsh_index: LSHIndex,
        expansion_depth: int = 1,
        min_similarity: float = 0.1,
        ngram_size: int = 3,
    ) -> None:
        self.lsh_index = lsh_index
        self.expansion_depth = expansion_depth
        self.min_similarity = min_similarity
        self.ngram_size = ngram_size
        # item_id -> metadata cache
        self._metadata: Dict[Hashable, Dict[str, Any]] = {}

    def set_metadata(self, item_id: Hashable, metadata: Dict[str, Any]) -> None:
        """Attach metadata to an item for enriched results."""
        self._metadata[item_id] = metadata

    def resolve(
        self,
        signal: Signal,
        *,
        top_k: int = 20,
        min_similarity: Optional[float] = None,
    ) -> List[SemanticResult]:
        """Find semantically similar items for a query signal.

        Uses the signal's shingle set to query the LSH index.  Results
        are ranked by estimated Jaccard similarity.

        Args:
            signal:         The query signal containing shingles.
            top_k:          Maximum number of results.
            min_similarity: Override instance-level minimum similarity.

        Returns:
            List of SemanticResult sorted by descending similarity.
        """
        threshold = min_similarity if min_similarity is not None else self.min_similarity

        if not signal.shingles:
            return []

        # Query LSH index - this does band-based candidate retrieval + re-ranking
        lsh_hits = self.lsh_index.query(
            signal.shingles,
            top_k=top_k * (1 + self.expansion_depth),
            min_similarity=threshold,
        )

        results: List[SemanticResult] = []
        for item_id, similarity in lsh_hits:
            results.append(SemanticResult(
                item_id=item_id,
                similarity=similarity,
                metadata=self._metadata.get(item_id),
            ))

        # Already sorted by LSH query, just truncate
        return results[:top_k]

    def resolve_text(
        self,
        text: str,
        *,
        top_k: int = 20,
        min_similarity: Optional[float] = None,
    ) -> List[SemanticResult]:
        """Convenience: resolve from raw text without a pre-built Signal.

        Shingles are computed on the fly.

        Args:
            text:           The query text.
            top_k:          Maximum number of results.
            min_similarity: Minimum similarity threshold.

        Returns:
            List of SemanticResult sorted by descending similarity.
        """
        threshold = min_similarity if min_similarity is not None else self.min_similarity

        char_shingles = ngram_shingles(text, n=self.ngram_size)
        w_shingles = word_shingles(text, n=2)
        all_shingles = char_shingles | w_shingles

        if not all_shingles:
            return []

        lsh_hits = self.lsh_index.query(
            all_shingles,
            top_k=top_k,
            min_similarity=threshold,
        )

        results: List[SemanticResult] = []
        for item_id, similarity in lsh_hits:
            results.append(SemanticResult(
                item_id=item_id,
                similarity=similarity,
                metadata=self._metadata.get(item_id),
            ))

        return results

    def find_clusters(
        self,
        signal: Signal,
    ) -> List[SemanticCluster]:
        """Identify the LSH clusters that a signal falls into.

        Examines each band of the LSH index and collects clusters
        (buckets) where the signal's shingle set matches.

        Args:
            signal: The query signal.

        Returns:
            List of SemanticCluster objects, one per matching bucket.
        """
        if not signal.minhash_signature:
            return []

        sig = list(signal.minhash_signature)
        clusters: List[SemanticCluster] = []

        for band_idx in range(self.lsh_index.bands):
            start = band_idx * self.lsh_index.rows
            band = tuple(sig[start: start + self.lsh_index.rows])
            band_hash = hash(band)

            bucket = self.lsh_index._buckets[band_idx].get(band_hash)
            if bucket:
                cluster = SemanticCluster(cluster_id=(band_idx, band_hash))
                for member_id in bucket:
                    cluster.add(member_id)
                clusters.append(cluster)

        return clusters

    def expand_neighbors(
        self,
        item_id: Hashable,
        *,
        top_k: int = 10,
        min_similarity: Optional[float] = None,
    ) -> List[SemanticResult]:
        """Find semantic neighbors of an already-indexed item.

        Args:
            item_id:        ID of the item to find neighbors for.
            top_k:          Maximum number of neighbors.
            min_similarity: Minimum similarity threshold.

        Returns:
            List of SemanticResult representing the nearest neighbors.

        Raises:
            KeyError: If item_id is not in the LSH index.
        """
        threshold = min_similarity if min_similarity is not None else self.min_similarity

        lsh_hits = self.lsh_index.query_by_id(
            item_id,
            top_k=top_k,
            min_similarity=threshold,
        )

        results: List[SemanticResult] = []
        for neighbor_id, similarity in lsh_hits:
            results.append(SemanticResult(
                item_id=neighbor_id,
                similarity=similarity,
                metadata=self._metadata.get(neighbor_id),
            ))

        return results
