"""Routing helpers for selecting the right index or shard for a request."""
from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Mapping, Optional, Protocol, Sequence

import numpy as np


class _RoutableIndex(Protocol):
    @property
    def ids(self) -> Sequence[str]: ...

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

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

    def add(self, episode_id: str, vec: np.ndarray) -> None: ...

    def remove(self, episode_id: str) -> None: ...

    def search(self,
               query: np.ndarray,
               k: int) -> Sequence[tuple[str,
                                         float]]: ...

    def clear(self) -> None: ...


@dataclass
class IndexRoute:
    namespace: Optional[str] = None
    tenant: Optional[str] = None
    encoder_id: Optional[str] = None


class IndexRouter:
    """Simple in-process router for vector indices.

    The router keeps namespace and encoder specific mappings while
    falling back to a default index. It is intentionally minimal so it can
    later be replaced by a distributed coordinator without changing call sites.
    """

    def __init__(
        self,
        default_index: Optional[_RoutableIndex] = None,
        namespace_mapping: Optional[Mapping[str, _RoutableIndex]] = None,
        encoder_mapping: Optional[Mapping[str, _RoutableIndex]] = None,
    ) -> None:
        self.default_index = default_index
        self.namespace_mapping: Dict[str, _RoutableIndex] = dict(
            namespace_mapping or {})
        self.encoder_mapping: Dict[str, _RoutableIndex] = dict(
            encoder_mapping or {})

    def route(
            self,
            *,
            namespace: Optional[str] = None,
            encoder_id: Optional[str] = None) -> Optional[_RoutableIndex]:
        if encoder_id and encoder_id in self.encoder_mapping:
            return self.encoder_mapping[encoder_id]
        if namespace and namespace in self.namespace_mapping:
            return self.namespace_mapping[namespace]
        return self.default_index

    def set_default(self, index: _RoutableIndex) -> None:
        self.default_index = index

    def register_namespace(
            self,
            namespace: str,
            index: _RoutableIndex) -> None:
        self.namespace_mapping[namespace] = index

    def register_encoder(self, encoder_id: str, index: _RoutableIndex) -> None:
        self.encoder_mapping[encoder_id] = index
