"""Sharding helpers for vector backends.

These helpers provide a lightweight abstraction around the existing
:class:`PGVectorStore` and :class:`FaissIndex` implementations so that the rest
of the codebase can treat a cluster of shards as a single store. The initial
version focuses on static shard maps and in-process fan-out, but the API is
structured so we can later move the routing and fan-out to remote workers.
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple

import numpy as np

from hieromem_core.episode import Episode
from hieromem_core.persistence.faiss_index import FaissIndex
from hieromem_core.persistence.pgvector_store import PGVectorStore

ShardKeyFn = Callable[[Mapping[str, str] | None], int]


@dataclass(frozen=True)
class ShardConfig:
    """Static configuration for a shard."""

    shard_id: int
    table_name: Optional[str] = None
    index_path: Optional[str] = None


def default_hash_shard(num_shards: int) -> ShardKeyFn:
    def _resolver(meta: Mapping[str, str] | None) -> int:
        if not meta:
            return 0
        key = meta.get("tenant") or meta.get(
            "namespace") or meta.get("shard_key")
        if key is None:
            return 0
        return hash(str(key)) % max(1, num_shards)

    return _resolver


class ShardedPGVectorStore:
    """Fan-out wrapper over multiple :class:`PGVectorStore` shards."""

    def __init__(
        self,
        shards: Mapping[int, PGVectorStore],
        *,
        shard_fn: Optional[ShardKeyFn] = None,
    ) -> None:
        self.shards: Dict[int, PGVectorStore] = dict(shards)
        self.shard_fn = shard_fn or default_hash_shard(len(self.shards) or 1)

    def _select_shard(self, meta: Mapping[str, str] | None) -> PGVectorStore:
        if not self.shards:
            raise ValueError("No shards configured for ShardedPGVectorStore")
        shard_ids = sorted(self.shards)
        sid = shard_ids[self.shard_fn(meta) % len(shard_ids)]
        return self.shards[sid]

    def save_episode(self, ep: Episode) -> None:
        shard = self._select_shard(getattr(ep, "meta", None))
        shard.save_episode(ep)

    def save_many(self, episodes: Iterable[Episode]) -> None:
        buckets: Dict[int, List[Episode]] = {}
        shard_ids = sorted(self.shards)
        for ep in episodes:
            sid = shard_ids[self.shard_fn(
                getattr(ep, "meta", None)) % len(shard_ids)]
            buckets.setdefault(sid, []).append(ep)
        for sid, eps in buckets.items():
            self.shards[sid].save_many(eps)

    def search_similar(self, query_vec: np.ndarray,
                       k: int = 5) -> List[Tuple[str, float]]:
        if not self.shards:
            return []
        results: List[Tuple[str, float]] = []
        for shard in self.shards.values():
            results.extend(shard.search_similar(query_vec, k))
        results.sort(key=lambda x: x[1], reverse=True)
        return results[:k]

    def close(self) -> None:
        for shard in self.shards.values():
            shard.close()


class ShardedFaissIndex:
    """Sharded FAISS wrapper that merges top-k across shards."""

    def __init__(
        self,
        shards: Mapping[int, FaissIndex],
        *,
        shard_fn: Optional[ShardKeyFn] = None,
    ) -> None:
        self.shards: Dict[int, FaissIndex] = dict(shards)
        self.shard_fn = shard_fn or default_hash_shard(len(self.shards) or 1)

    @property
    def ids(self) -> List[str]:
        all_ids: List[str] = []
        for shard in self.shards.values():
            all_ids.extend(list(getattr(shard, "ids", [])))
        return all_ids

    def _select_shard(self, meta: Mapping[str, str] | None) -> FaissIndex:
        if not self.shards:
            raise ValueError("No shards configured for ShardedFaissIndex")
        shard_ids = sorted(self.shards)
        sid = shard_ids[self.shard_fn(meta) % len(shard_ids)]
        return self.shards[sid]

    def add_embeddings(self,
                       ids: Sequence[str],
                       matrix: np.ndarray,
                       *,
                       meta: Mapping[str,
                                     str] | None = None) -> None:
        if not ids:
            return
        shard = self._select_shard(meta)
        shard.add_embeddings(ids, matrix)

    def add_many(self, ids: Sequence[str], matrix: np.ndarray) -> None:
        if not ids:
            return
        self.add_embeddings(ids, matrix)

    def add(self, episode_id: str, vec: np.ndarray, *, meta: Mapping[str, str] | None = None) -> None:
        shard = self._select_shard(meta)
        shard.add(episode_id, vec)

    def remove(self, episode_id: str) -> None:
        for shard in self.shards.values():
            shard.remove(episode_id)

    def clear(self) -> None:
        for shard in self.shards.values():
            shard.clear()

    def search(self, query: np.ndarray, k: int = 5) -> List[Tuple[str, float]]:
        merged: List[Tuple[str, float]] = []
        for shard in self.shards.values():
            merged.extend(shard.search(query, k))
        merged.sort(key=lambda x: x[1], reverse=True)
        return merged[:k]
