"""
RBSP Distributed Sharding.

Provides consistent-hash-based sharding of documents across multiple
local Shard partitions, each wrapping its own VectorStore + InvertedIndex.
Scatter-gather query coordination merges results from all shards.

Components:
    - ConsistencyLevel:  Enum for read/write consistency.
    - ShardConfig:       Configuration for the sharded index.
    - ShardRouter:       Routes doc_ids to shards via jump consistent hash.
    - Shard:             A single partition wrapping InvertedIndex + VectorStore.
    - ReplicaManager:    Manages read replicas with health tracking.
    - QueryCoordinator:  Scatter-gather with timeout and partial results.
    - ShardedIndex:      Main entry point for distributed operations.

Thread-safe via ``threading.RLock``.  Zero external dependencies.
"""

from __future__ import annotations

import heapq
import threading
import time
from enum import Enum
from typing import Any, Dict, List, Optional, Sequence, Set, Tuple


# ---------------------------------------------------------------------------
# Consistency Level
# ---------------------------------------------------------------------------

class ConsistencyLevel(Enum):
    """Read/write consistency levels for replica operations."""

    ONE = 1       # Fastest: succeed after 1 replica responds
    QUORUM = 2    # Majority of replicas must respond
    ALL = 3       # Every replica must respond


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

class ShardConfig:
    """Configuration for the sharded index.

    Args:
        num_shards:         Number of shard partitions (power of 2 recommended).
        replication_factor: Number of copies of each shard (1 = no replication).
        consistency_level:  Read/write consistency requirement.
        scatter_timeout:    Seconds to wait for shard responses during search.
        allow_partial:      If True, return partial results on shard timeout.
    """

    __slots__ = (
        "num_shards",
        "replication_factor",
        "consistency_level",
        "scatter_timeout",
        "allow_partial",
    )

    def __init__(
        self,
        *,
        num_shards: int = 4,
        replication_factor: int = 1,
        consistency_level: ConsistencyLevel = ConsistencyLevel.ONE,
        scatter_timeout: float = 5.0,
        allow_partial: bool = True,
    ) -> None:
        if num_shards < 1:
            raise ValueError("num_shards must be >= 1")
        if replication_factor < 1:
            raise ValueError("replication_factor must be >= 1")
        if scatter_timeout <= 0:
            raise ValueError("scatter_timeout must be > 0")
        self.num_shards = num_shards
        self.replication_factor = replication_factor
        self.consistency_level = consistency_level
        self.scatter_timeout = scatter_timeout
        self.allow_partial = allow_partial

    def __repr__(self) -> str:
        return (
            f"ShardConfig(num_shards={self.num_shards}, "
            f"replication_factor={self.replication_factor}, "
            f"consistency={self.consistency_level.name})"
        )


# ---------------------------------------------------------------------------
# Jump Consistent Hash
# ---------------------------------------------------------------------------

def jump_consistent_hash(key: int, num_buckets: int) -> int:
    """Map *key* to a bucket in [0, num_buckets) using jump consistent hash.

    Properties:
        - O(ln(num_buckets)) time
        - Uniform distribution
        - Minimal reassignment when num_buckets changes

    Reference: Lamping & Veach, "A Fast, Minimal Memory, Consistent Hash
    Algorithm" (Google, 2014).

    Args:
        key:         Non-negative integer key.
        num_buckets: Number of buckets (must be >= 1).

    Returns:
        Bucket index in [0, num_buckets).

    Raises:
        ValueError: If num_buckets < 1.
    """
    if num_buckets <= 0:
        raise ValueError("num_buckets must be >= 1")
    if num_buckets == 1:
        return 0

    b: int = -1
    j: int = 0
    seed = key
    while j < num_buckets:
        b = j
        seed = ((seed * 2862933555777941757) + 1) & 0xFFFFFFFFFFFFFFFF
        j = int((b + 1) * (1 << 31) / ((seed >> 33) + 1))
    return b


def _doc_id_to_int(doc_id: str) -> int:
    """Convert a doc_id string to a non-negative integer for hashing.

    Uses Python's built-in hash with masking to ensure non-negative.
    """
    return hash(doc_id) & 0xFFFFFFFFFFFFFFFF


# ---------------------------------------------------------------------------
# Shard Router
# ---------------------------------------------------------------------------

class ShardRouter:
    """Routes document IDs to shard indices using jump consistent hash.

    Thread-safe.  The number of shards can be updated for rebalancing.

    Args:
        num_shards: Initial number of shards.
    """

    def __init__(self, num_shards: int) -> None:
        if num_shards < 1:
            raise ValueError("num_shards must be >= 1")
        self._num_shards = num_shards
        self._lock = threading.Lock()

    @property
    def num_shards(self) -> int:
        """Number of shard partitions."""
        with self._lock:
            return self._num_shards

    @num_shards.setter
    def num_shards(self, value: int) -> None:
        if value < 1:
            raise ValueError("num_shards must be >= 1")
        with self._lock:
            self._num_shards = value

    def route(self, doc_id: str) -> int:
        """Return the shard index for *doc_id*.

        Args:
            doc_id: Document identifier string.

        Returns:
            Shard index in [0, num_shards).
        """
        key = _doc_id_to_int(doc_id)
        with self._lock:
            return jump_consistent_hash(key, self._num_shards)

    def route_with_replicas(self, doc_id: str, replication_factor: int) -> List[int]:
        """Return primary + replica shard indices for *doc_id*.

        Uses successive hash offsets to pick distinct replicas.

        Args:
            doc_id:             Document identifier string.
            replication_factor: Total number of copies (including primary).

        Returns:
            List of distinct shard indices, length = min(replication_factor, num_shards).
        """
        key = _doc_id_to_int(doc_id)
        with self._lock:
            n = self._num_shards
        rf = min(replication_factor, n)
        shards: List[int] = []
        seen: Set[int] = set()
        offset = 0
        while len(shards) < rf:
            shard_id = jump_consistent_hash(
                (key + offset) & 0xFFFFFFFFFFFFFFFF, n
            )
            if shard_id not in seen:
                seen.add(shard_id)
                shards.append(shard_id)
            offset += 1
            # Safety valve
            if offset > n * 10:
                break
        return shards

    def __repr__(self) -> str:
        return f"ShardRouter(num_shards={self._num_shards})"


# ---------------------------------------------------------------------------
# Shard -- wraps local VectorStore + InvertedIndex for one partition
# ---------------------------------------------------------------------------

class Shard:
    """A single shard partition holding a subset of indexed documents.

    Wraps a VectorStore (dense vectors) and an InvertedIndex (sparse tokens)
    for the documents assigned to this shard.

    Args:
        shard_id:   Unique identifier for this shard.
        vector_dim: Dimensionality for VectorStore (None to infer on first
                    insert).
    """

    def __init__(self, shard_id: int, vector_dim: Optional[int] = None) -> None:
        from rbsp.indexer.inverted import InvertedIndex
        from rbsp.indexer.vector_store import VectorStore

        self.shard_id = shard_id
        self._inverted = InvertedIndex()
        self._vector_store = VectorStore(dim=vector_dim)
        self._doc_ids: Set[str] = set()
        self._lock = threading.RLock()

    @property
    def inverted(self) -> Any:
        """The shard's InvertedIndex."""
        return self._inverted

    @property
    def vector_store(self) -> Any:
        """The shard's VectorStore."""
        return self._vector_store

    @property
    def doc_count(self) -> int:
        """Number of documents in this shard."""
        with self._lock:
            return len(self._doc_ids)

    def insert(
        self,
        doc_id: str,
        tokens: Optional[List[str]] = None,
        vector: Optional[Sequence[float]] = None,
    ) -> None:
        """Insert a document into this shard.

        Args:
            doc_id: Unique document identifier.
            tokens: Token list for inverted index (optional).
            vector: Dense vector for vector store (optional).
        """
        with self._lock:
            self._doc_ids.add(doc_id)
            if tokens is not None:
                self._inverted.add_document(doc_id, tokens)
            if vector is not None:
                self._vector_store.insert(doc_id, vector)

    def remove(self, doc_id: str) -> bool:
        """Remove a document from this shard.

        Returns:
            True if the document was found and removed.
        """
        with self._lock:
            if doc_id not in self._doc_ids:
                return False
            self._doc_ids.discard(doc_id)
            self._inverted.remove_document(doc_id)
            self._vector_store.remove(doc_id)
            return True

    def contains(self, doc_id: str) -> bool:
        """Check if this shard contains a document."""
        with self._lock:
            return doc_id in self._doc_ids

    def all_doc_ids(self) -> Set[str]:
        """Return all document IDs in this shard (copy)."""
        with self._lock:
            return set(self._doc_ids)

    def search_tokens(self, term: str) -> Set[str]:
        """Lookup a single term in the inverted index.

        Args:
            term: Token to look up.

        Returns:
            Set of doc_ids containing the term.
        """
        return self._inverted.lookup(term)

    def search_vector(
        self,
        query: Sequence[float],
        k: int = 10,
        mode: str = "hnsw",
    ) -> List[Tuple[str, float]]:
        """Search the vector store.

        Args:
            query: Query vector.
            k:     Number of results.
            mode:  VectorStore search mode.

        Returns:
            List of (doc_id, distance) pairs sorted by ascending distance.
        """
        return self._vector_store.search(query, k, mode=mode)

    def search_boolean(self, query_str: str) -> Set[str]:
        """Evaluate a boolean query string on the inverted index.

        Args:
            query_str: Boolean query (e.g. "python AND hash").

        Returns:
            Set of matching doc_ids.
        """
        return self._inverted.search(query_str)

    def __len__(self) -> int:
        return self.doc_count

    def __contains__(self, doc_id: str) -> bool:
        return self.contains(doc_id)

    def __repr__(self) -> str:
        return f"Shard(id={self.shard_id}, docs={self.doc_count})"


# ---------------------------------------------------------------------------
# Replica Manager
# ---------------------------------------------------------------------------

class ReplicaManager:
    """Manages read replicas for shard partitions.

    Each shard can have multiple replicas.  Writes go to all replicas;
    reads are served from the first healthy replica (with fallback).

    Args:
        config: ShardConfig governing replication behavior.
    """

    def __init__(self, config: ShardConfig) -> None:
        self._config = config
        # _replicas[shard_id] = list of Shard replicas (index 0 = primary)
        self._replicas: Dict[int, List[Shard]] = {}
        # (shard_id, replica_idx) -> healthy
        self._health: Dict[Tuple[int, int], bool] = {}
        self._lock = threading.RLock()

    def register_shard(self, shard: Shard) -> None:
        """Register a shard as primary (replica index 0).

        Args:
            shard: The primary shard to register.
        """
        with self._lock:
            sid = shard.shard_id
            if sid not in self._replicas:
                self._replicas[sid] = []
            if not self._replicas[sid]:
                self._replicas[sid].append(shard)
            else:
                self._replicas[sid][0] = shard
            self._health[(sid, 0)] = True

    def add_replica(self, shard_id: int, replica: Shard) -> int:
        """Add a read replica for *shard_id*.

        Args:
            shard_id: Primary shard to add a replica for.
            replica:  The replica Shard instance.

        Returns:
            The replica index.
        """
        with self._lock:
            if shard_id not in self._replicas:
                self._replicas[shard_id] = []
            idx = len(self._replicas[shard_id])
            self._replicas[shard_id].append(replica)
            self._health[(shard_id, idx)] = True
            return idx

    def get_read_shard(self, shard_id: int) -> Optional[Shard]:
        """Get a healthy replica for reading.

        Returns the first healthy replica, or None if all are down.

        Args:
            shard_id: Shard partition to read from.

        Returns:
            A healthy Shard, or None.
        """
        with self._lock:
            replicas = self._replicas.get(shard_id, [])
            for idx, replica in enumerate(replicas):
                if self._health.get((shard_id, idx), False):
                    return replica
            return None

    def get_all_replicas(self, shard_id: int) -> List[Shard]:
        """Get all replicas (including primary) for a shard.

        Args:
            shard_id: Shard partition.

        Returns:
            List of all Shard replicas.
        """
        with self._lock:
            return list(self._replicas.get(shard_id, []))

    def get_write_replicas(
        self, shard_id: int, consistency: ConsistencyLevel
    ) -> List[Shard]:
        """Get replicas that must be written to for the given consistency level.

        Args:
            shard_id:    Target shard.
            consistency: Required consistency level.

        Returns:
            List of Shard replicas to write to.
        """
        with self._lock:
            replicas = self._replicas.get(shard_id, [])
            healthy = [
                (idx, r)
                for idx, r in enumerate(replicas)
                if self._health.get((shard_id, idx), False)
            ]
            if consistency == ConsistencyLevel.ALL:
                return [r for _, r in healthy]
            elif consistency == ConsistencyLevel.QUORUM:
                quorum = len(replicas) // 2 + 1
                return [r for _, r in healthy[:quorum]]
            else:  # ONE
                return [r for _, r in healthy[:1]]

    def mark_healthy(self, shard_id: int, replica_idx: int) -> None:
        """Mark a replica as healthy.

        Args:
            shard_id:    Shard partition.
            replica_idx: Replica index within the partition.
        """
        with self._lock:
            self._health[(shard_id, replica_idx)] = True

    def mark_unhealthy(self, shard_id: int, replica_idx: int) -> None:
        """Mark a replica as unhealthy.

        Args:
            shard_id:    Shard partition.
            replica_idx: Replica index within the partition.
        """
        with self._lock:
            self._health[(shard_id, replica_idx)] = False

    def replica_count(self, shard_id: int) -> int:
        """Return the number of replicas for *shard_id*.

        Args:
            shard_id: Shard partition.

        Returns:
            Number of replicas (including primary).
        """
        with self._lock:
            return len(self._replicas.get(shard_id, []))

    def __repr__(self) -> str:
        with self._lock:
            total = sum(len(r) for r in self._replicas.values())
        return f"ReplicaManager(total_replicas={total})"


# ---------------------------------------------------------------------------
# Query Coordinator
# ---------------------------------------------------------------------------

class QueryCoordinator:
    """Coordinates scatter-gather queries across shards.

    Fans out queries to all shards in parallel (using threads), then
    merges results.  Supports timeouts and partial results on failure.

    Args:
        config: ShardConfig governing timeout and partial result behavior.
    """

    def __init__(self, config: ShardConfig) -> None:
        self._config = config

    def scatter_gather_vector(
        self,
        shards: List[Shard],
        query: Sequence[float],
        k: int = 10,
        mode: str = "hnsw",
    ) -> List[Tuple[str, float]]:
        """Scatter a vector search to all shards and gather/merge results.

        Each shard returns its local top-k; this method merges and returns
        the global top-k by ascending distance.

        Args:
            shards: List of Shard instances to query.
            query:  Query vector.
            k:      Number of global results.
            mode:   VectorStore search mode.

        Returns:
            Global top-k results as (doc_id, distance) sorted ascending.
        """
        if not shards:
            return []

        results_per_shard: List[List[Tuple[str, float]]] = [
            [] for _ in shards
        ]
        errors: List[Optional[Exception]] = [None] * len(shards)

        def _query_shard(idx: int, shard: Shard) -> None:
            try:
                results_per_shard[idx] = shard.search_vector(query, k, mode)
            except Exception as e:
                errors[idx] = e

        threads: List[threading.Thread] = []
        for i, shard in enumerate(shards):
            t = threading.Thread(
                target=_query_shard, args=(i, shard), daemon=True
            )
            t.start()
            threads.append(t)

        # Wait for all shards with timeout
        deadline = time.monotonic() + self._config.scatter_timeout
        for t in threads:
            remaining = deadline - time.monotonic()
            if remaining > 0:
                t.join(timeout=remaining)

        # Merge: collect all shard results
        all_results: List[Tuple[str, float]] = []
        for shard_results in results_per_shard:
            all_results.extend(shard_results)

        # Return global top-k by smallest distance
        return heapq.nsmallest(k, all_results, key=lambda x: x[1])

    def scatter_gather_tokens(
        self,
        shards: List[Shard],
        term: str,
    ) -> Set[str]:
        """Scatter a single-term lookup to all shards and union results.

        Args:
            shards: List of Shard instances to query.
            term:   Token to look up.

        Returns:
            Union of all doc_ids matching the term across all shards.
        """
        if not shards:
            return set()

        results_per_shard: List[Set[str]] = [set() for _ in shards]

        def _query_shard(idx: int, shard: Shard) -> None:
            try:
                results_per_shard[idx] = shard.search_tokens(term)
            except Exception:
                pass

        threads: List[threading.Thread] = []
        for i, shard in enumerate(shards):
            t = threading.Thread(
                target=_query_shard, args=(i, shard), daemon=True
            )
            t.start()
            threads.append(t)

        deadline = time.monotonic() + self._config.scatter_timeout
        for t in threads:
            remaining = deadline - time.monotonic()
            if remaining > 0:
                t.join(timeout=remaining)

        result: Set[str] = set()
        for s in results_per_shard:
            result.update(s)
        return result

    def scatter_gather_boolean(
        self,
        shards: List[Shard],
        query_str: str,
    ) -> Set[str]:
        """Scatter a boolean query to all shards and union results.

        Args:
            shards:    Shard instances to query.
            query_str: Boolean query string (e.g. "python AND hash").

        Returns:
            Union of matching doc_ids across all shards.
        """
        if not shards:
            return set()

        results_per_shard: List[Set[str]] = [set() for _ in shards]

        def _query_shard(idx: int, shard: Shard) -> None:
            try:
                results_per_shard[idx] = shard.search_boolean(query_str)
            except Exception:
                pass

        threads: List[threading.Thread] = []
        for i, shard in enumerate(shards):
            t = threading.Thread(
                target=_query_shard, args=(i, shard), daemon=True
            )
            t.start()
            threads.append(t)

        deadline = time.monotonic() + self._config.scatter_timeout
        for t in threads:
            remaining = deadline - time.monotonic()
            if remaining > 0:
                t.join(timeout=remaining)

        result: Set[str] = set()
        for s in results_per_shard:
            result.update(s)
        return result


# ---------------------------------------------------------------------------
# Sharded Index -- the main entry point
# ---------------------------------------------------------------------------

class ShardedIndex:
    """Manages multiple Shard instances with consistent-hash routing.

    Provides insert, search, and rebalance operations that transparently
    route to the correct shard(s).

    Args:
        config:     Sharding configuration.
        vector_dim: Vector dimensionality (optional, inferred on first insert).

    Usage::

        idx = ShardedIndex(ShardConfig(num_shards=4))
        idx.insert("doc1", tokens=["python", "search", "engine"])
        results = idx.search_tokens("python")
    """

    def __init__(
        self,
        config: Optional[ShardConfig] = None,
        vector_dim: Optional[int] = None,
    ) -> None:
        self._config = config or ShardConfig()
        self._vector_dim = vector_dim
        self._router = ShardRouter(self._config.num_shards)
        self._shards: Dict[int, Shard] = {}
        self._replica_mgr = ReplicaManager(self._config)
        self._coordinator = QueryCoordinator(self._config)
        self._lock = threading.RLock()

        # Initialize shards
        for i in range(self._config.num_shards):
            shard = Shard(shard_id=i, vector_dim=vector_dim)
            self._shards[i] = shard
            self._replica_mgr.register_shard(shard)

            # Create additional replicas
            for _ in range(self._config.replication_factor - 1):
                replica = Shard(shard_id=i, vector_dim=vector_dim)
                self._replica_mgr.add_replica(i, replica)

    @property
    def config(self) -> ShardConfig:
        """The shard configuration."""
        return self._config

    @property
    def num_shards(self) -> int:
        """Number of shard partitions."""
        return self._router.num_shards

    @property
    def total_docs(self) -> int:
        """Total documents across all primary shards."""
        with self._lock:
            return sum(s.doc_count for s in self._shards.values())

    def shard_sizes(self) -> Dict[int, int]:
        """Return {shard_id: doc_count} for all primary shards."""
        with self._lock:
            return {sid: s.doc_count for sid, s in self._shards.items()}

    # -- Insert ---------------------------------------------------------------

    def insert(
        self,
        doc_id: str,
        tokens: Optional[List[str]] = None,
        vector: Optional[Sequence[float]] = None,
    ) -> int:
        """Insert a document, routing to the correct shard.

        Also writes to replicas according to the consistency level.

        Args:
            doc_id: Unique document identifier.
            tokens: Token list for inverted index.
            vector: Dense vector for vector store.

        Returns:
            The shard_id the document was routed to.
        """
        shard_id = self._router.route(doc_id)

        with self._lock:
            # Write to primary
            primary = self._shards[shard_id]
            primary.insert(doc_id, tokens=tokens, vector=vector)

            # Write to replicas
            write_targets = self._replica_mgr.get_write_replicas(
                shard_id, self._config.consistency_level
            )
            for replica in write_targets:
                if replica is not primary:
                    replica.insert(doc_id, tokens=tokens, vector=vector)

        return shard_id

    def batch_insert(
        self,
        doc_ids: List[str],
        tokens_list: Optional[List[Optional[List[str]]]] = None,
        vectors: Optional[List[Optional[Sequence[float]]]] = None,
    ) -> Dict[int, int]:
        """Batch-insert multiple documents.

        Args:
            doc_ids:     List of document identifiers.
            tokens_list: Parallel list of token lists (or None entries).
            vectors:     Parallel list of vectors (or None entries).

        Returns:
            Dict mapping shard_id -> number of documents inserted.
        """
        counts: Dict[int, int] = {}
        n = len(doc_ids)
        for i in range(n):
            toks = (
                tokens_list[i]
                if tokens_list and i < len(tokens_list)
                else None
            )
            vec = vectors[i] if vectors and i < len(vectors) else None
            sid = self.insert(doc_ids[i], tokens=toks, vector=vec)
            counts[sid] = counts.get(sid, 0) + 1
        return counts

    # -- Remove ---------------------------------------------------------------

    def remove(self, doc_id: str) -> bool:
        """Remove a document from its shard and replicas.

        Args:
            doc_id: Document to remove.

        Returns:
            True if the document was found and removed.
        """
        shard_id = self._router.route(doc_id)
        with self._lock:
            primary = self._shards[shard_id]
            removed = primary.remove(doc_id)
            # Also remove from replicas
            for replica in self._replica_mgr.get_all_replicas(shard_id):
                if replica is not primary:
                    replica.remove(doc_id)
            return removed

    # -- Search ---------------------------------------------------------------

    def search_vector(
        self,
        query: Sequence[float],
        k: int = 10,
        mode: str = "hnsw",
    ) -> List[Tuple[str, float]]:
        """Scatter-gather vector search across all shards.

        Args:
            query: Query vector.
            k:     Number of results.
            mode:  VectorStore search mode.

        Returns:
            Global top-k (doc_id, distance) sorted by ascending distance.
        """
        read_shards = self._get_read_shards()
        return self._coordinator.scatter_gather_vector(
            read_shards, query, k, mode
        )

    def search_tokens(self, term: str) -> Set[str]:
        """Single-term lookup across all shards.

        Args:
            term: Token to look up.

        Returns:
            Union of matching doc_ids.
        """
        read_shards = self._get_read_shards()
        return self._coordinator.scatter_gather_tokens(read_shards, term)

    def search_boolean(self, query_str: str) -> Set[str]:
        """Boolean query across all shards.

        Args:
            query_str: Boolean query string (e.g. "python AND hash").

        Returns:
            Union of matching doc_ids.
        """
        read_shards = self._get_read_shards()
        return self._coordinator.scatter_gather_boolean(
            read_shards, query_str
        )

    # -- Rebalance ------------------------------------------------------------

    def rebalance(self, new_num_shards: int) -> Dict[str, int]:
        """Redistribute all documents across *new_num_shards* shards.

        Creates new shards, re-routes every document via the updated router,
        and migrates data.  Old shards are discarded.

        Args:
            new_num_shards: Target number of shards.

        Returns:
            Migration stats with keys ``"moved"``, ``"total"``,
            ``"old_shards"``, ``"new_shards"``.

        Raises:
            ValueError: If new_num_shards < 1.
        """
        if new_num_shards < 1:
            raise ValueError("new_num_shards must be >= 1")

        with self._lock:
            old_num = self._router.num_shards
            if new_num_shards == old_num:
                return {
                    "moved": 0,
                    "total": self.total_docs,
                    "old_shards": old_num,
                    "new_shards": new_num_shards,
                }

            # Collect all documents from current shards
            all_docs: List[
                Tuple[str, Optional[List[str]], Optional[List[float]]]
            ] = []
            for shard in self._shards.values():
                for doc_id in shard.all_doc_ids():
                    vec = shard.vector_store.get_vector(doc_id)
                    # Access tokens from InvertedIndex internal state
                    tokens_set = shard.inverted._doc_tokens.get(
                        doc_id, set()
                    )
                    tokens = list(tokens_set) if tokens_set else None
                    all_docs.append((doc_id, tokens, vec))

            # Create new shards
            new_shards: Dict[int, Shard] = {}
            new_replica_mgr = ReplicaManager(self._config)
            for i in range(new_num_shards):
                s = Shard(shard_id=i, vector_dim=self._vector_dim)
                new_shards[i] = s
                new_replica_mgr.register_shard(s)
                for _ in range(self._config.replication_factor - 1):
                    replica = Shard(
                        shard_id=i, vector_dim=self._vector_dim
                    )
                    new_replica_mgr.add_replica(i, replica)

            # Update router
            self._router.num_shards = new_num_shards

            # Re-insert all documents
            moved = 0
            for doc_id, tokens, vec in all_docs:
                new_shard_id = self._router.route(doc_id)
                new_shards[new_shard_id].insert(
                    doc_id, tokens=tokens, vector=vec
                )
                # Write to replicas
                for replica in new_replica_mgr.get_all_replicas(
                    new_shard_id
                ):
                    if replica is not new_shards[new_shard_id]:
                        replica.insert(
                            doc_id, tokens=tokens, vector=vec
                        )
                moved += 1

            # Swap
            self._shards = new_shards
            self._replica_mgr = new_replica_mgr

            return {
                "moved": moved,
                "total": len(all_docs),
                "old_shards": old_num,
                "new_shards": new_num_shards,
            }

    # -- Internals ------------------------------------------------------------

    def _get_read_shards(self) -> List[Shard]:
        """Get one healthy shard per partition for reading."""
        shards: List[Shard] = []
        with self._lock:
            for sid in range(self._router.num_shards):
                shard = self._replica_mgr.get_read_shard(sid)
                if shard is None:
                    # Fallback to primary even if marked unhealthy
                    shard = self._shards.get(sid)
                if shard is not None:
                    shards.append(shard)
        return shards

    def get_shard(self, shard_id: int) -> Optional[Shard]:
        """Get the primary shard by ID.

        Args:
            shard_id: Shard partition index.

        Returns:
            The primary Shard, or None if not found.
        """
        with self._lock:
            return self._shards.get(shard_id)

    def get_shard_for_doc(self, doc_id: str) -> int:
        """Return which shard a document would be routed to.

        Args:
            doc_id: Document identifier.

        Returns:
            Shard index in [0, num_shards).
        """
        return self._router.route(doc_id)

    def __repr__(self) -> str:
        return (
            f"ShardedIndex(shards={self.num_shards}, "
            f"docs={self.total_docs}, "
            f"replication={self._config.replication_factor})"
        )

    def __len__(self) -> int:
        return self.total_docs

    def __contains__(self, doc_id: str) -> bool:
        shard_id = self._router.route(doc_id)
        with self._lock:
            shard = self._shards.get(shard_id)
            if shard is None:
                return False
            return shard.contains(doc_id)
