from __future__ import annotations

from typing import Optional

import torch
import torch.nn.functional as F

try:  # pragma: no cover - optional dependency
    import faiss  # type: ignore

    _FAISS_AVAILABLE = True
except Exception:  # pragma: no cover - optional dependency
    faiss = None
    _FAISS_AVAILABLE = False


class ApproximateNearestNeighborIndex:
    """ANN helper that prefers FAISS when available and falls back to torch."""

    def __init__(self, *, use_faiss: Optional[bool] = None, use_gpu: bool = False):
        self._backend = _TorchExactIndex()
        target = _FAISS_AVAILABLE if use_faiss is None else use_faiss
        if target and _FAISS_AVAILABLE:
            self._backend = _FaissIndex(use_gpu=use_gpu)

    def rebuild(self, keys: torch.Tensor) -> None:
        self._backend.rebuild(keys)

    def search(self, queries: torch.Tensor, top_k: int):
        return self._backend.search(queries, top_k)


class _TorchExactIndex:
    def __init__(self) -> None:
        self._keys: Optional[torch.Tensor] = None

    def rebuild(self, keys: torch.Tensor) -> None:
        self._keys = keys.detach().cpu()

    def search(self, queries: torch.Tensor, top_k: int):
        if self._keys is None or self._keys.numel() == 0:
            empty_scores = torch.empty(queries.size(0), 0, dtype=torch.float32)
            empty_idx = torch.empty(queries.size(0), 0, dtype=torch.long)
            return empty_scores, empty_idx
        q = F.normalize(queries.detach().cpu().to(torch.float32), dim=-1)
        k = F.normalize(self._keys, dim=-1)
        scores = q @ k.t()
        return scores.topk(min(top_k, k.shape[0]), dim=-1)


class _FaissIndex:
    def __init__(self, *, use_gpu: bool = False) -> None:
        if not _FAISS_AVAILABLE:  # pragma: no cover - guard
            raise ImportError("faiss is not installed")
        self._use_gpu = use_gpu
        self._index = None
        self._dim = 0
        self._ntotal = 0
        self._gpu_resources = None

    def rebuild(self, keys: torch.Tensor) -> None:
        if keys.numel() == 0:
            self._index = None
            self._ntotal = 0
            return
        cpu_keys = F.normalize(keys.detach().cpu().to(torch.float32), dim=-1)
        self._dim = cpu_keys.shape[1]
        index = faiss.IndexFlatIP(self._dim)
        index.add(cpu_keys.numpy())
        if self._use_gpu:
            if not hasattr(faiss, "StandardGpuResources"):
                raise RuntimeError("faiss-gpu build required for GPU indexing")
            res = faiss.StandardGpuResources()
            index = faiss.index_cpu_to_gpu(res, 0, index)
            self._gpu_resources = res
        self._index = index
        self._ntotal = index.ntotal

    def search(self, queries: torch.Tensor, top_k: int):
        if self._index is None or self._ntotal == 0:
            empty_scores = torch.empty(queries.size(0), 0, dtype=torch.float32)
            empty_idx = torch.empty(queries.size(0), 0, dtype=torch.long)
            return empty_scores, empty_idx

        q = F.normalize(queries.detach().cpu().to(torch.float32), dim=-1)
        limit = min(top_k, self._ntotal)
        if limit <= 0:
            empty_scores = torch.empty(queries.size(0), 0, dtype=torch.float32)
            empty_idx = torch.empty(queries.size(0), 0, dtype=torch.long)
            return empty_scores, empty_idx
        scores_np, idx_np = self._index.search(q.numpy(), limit)
        scores = torch.from_numpy(scores_np)
        indices = torch.from_numpy(idx_np).to(dtype=torch.long)
        return scores, indices


class FaissIndex(ApproximateNearestNeighborIndex):
    """Explicit FAISS backend; raises if dependency is unavailable."""

    def __init__(self, *, use_gpu: bool = False):
        super().__init__(use_faiss=True, use_gpu=use_gpu)
        if not _FAISS_AVAILABLE:
            raise ImportError(
                "faiss is not installed; install hippotorch[faiss] to enable."
            )


__all__ = ["ApproximateNearestNeighborIndex", "FaissIndex", "_FAISS_AVAILABLE"]
