"""
RBSP Faceted Search and Aggregations.

Facets let users drill down into search results by categories such as
file type, directory, language, date range, etc.  Modeled after
Elasticsearch-style aggregations but built on pure Python with zero
external dependencies.

Components:
    - FacetValue:         A single bucket (value + count).
    - FacetResult:        Aggregation result for one field.
    - FacetCollector:     Collects and computes term/range facets.
    - HierarchicalFacet: Path-based faceting (e.g., directory trees).
"""

from __future__ import annotations

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


@dataclass(frozen=True, slots=True)
class FacetValue:
    """A single facet bucket.

    Attributes:
        value: The facet term or range label.
        count: Number of matching documents in this bucket.
    """

    value: str
    count: int


@dataclass(frozen=True, slots=True)
class FacetResult:
    """Aggregation result for a single facet field.

    Attributes:
        field:  The field name this facet was computed over.
        values: Top-N facet buckets ordered by count descending.
        total:  Total number of distinct values across matching docs.
    """

    field: str
    values: list[FacetValue]
    total: int


class FacetCollector:
    """Collects document facet metadata and computes aggregations.

    Thread-safe.  All mutations and reads are protected by a lock.

    Usage::

        collector = FacetCollector()
        collector.add_document("doc1", {"language": "python", "type": "file"})
        collector.add_document("doc2", {"language": "rust", "type": "file"})
        results = collector.collect({"doc1", "doc2"}, ["language"], top_n=5)
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        # field -> value -> set of doc_ids
        self._field_values: Dict[str, Dict[str, Set[str]]] = defaultdict(
            lambda: defaultdict(set)
        )
        # field -> doc_id -> numeric value (for range faceting)
        self._numeric_fields: Dict[str, Dict[str, float]] = defaultdict(dict)

    def add_document(self, doc_id: str, facet_fields: dict[str, str]) -> None:
        """Register a document's facet values.

        Args:
            doc_id:       Unique document identifier.
            facet_fields: Mapping of field name to the document's value
                          for that field (e.g., ``{"language": "python"}``).
        """
        with self._lock:
            for field_name, value in facet_fields.items():
                self._field_values[field_name][value].add(doc_id)
                # Attempt to store as numeric for range faceting
                try:
                    self._numeric_fields[field_name][doc_id] = float(value)
                except (ValueError, TypeError):
                    pass

    def collect(
        self,
        matching_docs: set[str],
        fields: list[str],
        top_n: int = 10,
    ) -> list[FacetResult]:
        """Compute term facets over a set of matching documents.

        For each requested field, counts how many matching documents have
        each distinct value and returns the top-N buckets by count.

        Args:
            matching_docs: The set of document IDs to aggregate over.
            fields:        Which facet fields to compute.
            top_n:         Maximum number of buckets per field.

        Returns:
            A list of FacetResult, one per requested field.
        """
        results: list[FacetResult] = []

        with self._lock:
            for field_name in fields:
                value_map = self._field_values.get(field_name)
                if value_map is None:
                    results.append(FacetResult(
                        field=field_name, values=[], total=0,
                    ))
                    continue

                # Count occurrences among matching docs
                counts: Dict[str, int] = {}
                for value, doc_ids in value_map.items():
                    overlap = len(doc_ids & matching_docs)
                    if overlap > 0:
                        counts[value] = overlap

                total_distinct = len(counts)

                # Top-N by count descending, stable on value ascending
                top = heapq.nlargest(
                    top_n,
                    counts.items(),
                    key=lambda item: (item[1], item[0]),
                )

                facet_values = [
                    FacetValue(value=val, count=cnt) for val, cnt in top
                ]

                results.append(FacetResult(
                    field=field_name,
                    values=facet_values,
                    total=total_distinct,
                ))

        return results

    def collect_range(
        self,
        matching_docs: set[str],
        field: str,
        ranges: list[tuple[float, float]],
    ) -> list[FacetValue]:
        """Compute numeric range facets over matching documents.

        Each range is a ``(low, high)`` tuple (inclusive on both ends).
        Documents whose numeric value for ``field`` falls within a range
        are counted in that bucket.

        Args:
            matching_docs: The set of document IDs to aggregate over.
            field:         The facet field (must have numeric values).
            ranges:        List of ``(low, high)`` range boundaries.

        Returns:
            A list of FacetValue, one per range, with the value formatted
            as ``"low-high"``.
        """
        results: list[FacetValue] = []

        with self._lock:
            numeric_map = self._numeric_fields.get(field, {})

            for low, high in ranges:
                count = 0
                for doc_id in matching_docs:
                    val = numeric_map.get(doc_id)
                    if val is not None and low <= val <= high:
                        count += 1
                label = f"{low}-{high}"
                results.append(FacetValue(value=label, count=count))

        return results


class HierarchicalFacet:
    """Path-based hierarchical faceting.

    Supports drill-down faceting over hierarchical paths such as directory
    trees.  Paths are split by a configurable separator and faceted at a
    requested depth level.

    Thread-safe.

    Usage::

        hf = HierarchicalFacet()
        hf.add_path("doc1", "src/rbsp/core/bloom.py")
        hf.add_path("doc2", "src/rbsp/core/lsh.py")
        hf.add_path("doc3", "src/rbsp/query/signal.py")
        buckets = hf.collect({"doc1", "doc2", "doc3"}, depth=2)
        # -> [FacetValue(value="src/rbsp", count=3)]
    """

    def __init__(self) -> None:
        self._lock = threading.Lock()
        # doc_id -> list of (path, separator) pairs
        self._paths: Dict[str, List[Tuple[str, str]]] = defaultdict(list)

    def add_path(
        self, doc_id: str, path: str, separator: str = "/"
    ) -> None:
        """Register a hierarchical path for a document.

        A document may have multiple paths (e.g., symlinks, aliases).

        Args:
            doc_id:    Unique document identifier.
            path:      The full hierarchical path.
            separator: Path component separator (default ``"/"``).
        """
        with self._lock:
            self._paths[doc_id].append((path, separator))

    def collect(
        self,
        matching_docs: set[str],
        depth: int = 1,
    ) -> list[FacetValue]:
        """Compute hierarchical facets at a given depth.

        Splits each document's path(s) by their separator and groups at
        the requested depth level.  Depth 1 gives the top-level component,
        depth 2 gives the first two components joined, etc.

        Args:
            matching_docs: The set of document IDs to aggregate over.
            depth:         How many path components deep to facet (>= 1).

        Returns:
            A list of FacetValue sorted by count descending, then by
            value ascending.
        """
        if depth < 1:
            raise ValueError(f"depth must be >= 1, got {depth}")

        counts: Dict[str, int] = {}

        with self._lock:
            for doc_id in matching_docs:
                entries = self._paths.get(doc_id)
                if entries is None:
                    continue

                # Track which prefixes this doc contributes to (avoid
                # double-counting when a doc has multiple paths with
                # the same prefix at this depth).
                seen_prefixes: Set[str] = set()

                for path, separator in entries:
                    parts = path.split(separator)
                    if len(parts) < depth:
                        continue
                    prefix = separator.join(parts[:depth])
                    if prefix not in seen_prefixes:
                        seen_prefixes.add(prefix)
                        counts[prefix] = counts.get(prefix, 0) + 1

        # Sort by count descending, then value ascending for stability
        sorted_items = sorted(
            counts.items(), key=lambda item: (-item[1], item[0])
        )
        return [FacetValue(value=val, count=cnt) for val, cnt in sorted_items]
