"""HNSW vector index for sub-linear similarity search."""

from __future__ import annotations

import base64
import heapq
import json
import math
import random
import threading
from pathlib import Path

import numpy as np


class HNSWIndex:
    """Hierarchical Navigable Small World graph for approximate nearest neighbors."""

    def __init__(
        self,
        dim: int = 384,
        M: int = 16,
        ef_construction: int = 200,
        ef_search: int = 50,
        max_elements: int = 1_000_000,
    ):
        self._dim = int(dim)
        self._M = max(2, int(M))
        self._M_max = self._M
        self._M_max0 = self._M * 2
        self._ef_construction = max(8, int(ef_construction))
        self._ef_search = max(8, int(ef_search))
        self._max_elements = int(max_elements)
        self._ml = 1.0 / math.log(self._M)

        self._max_level = 0
        self._entry_point = -1

        self._vectors: list[np.ndarray] = []
        self._ids: list[bytes] = []
        self._graphs: list[dict[int, list[int]]] = []
        self._id_to_node: dict[bytes, int] = {}

        self._lock = threading.RLock()

    @property
    def size(self) -> int:
        return len(self._vectors)

    @property
    def dimension(self) -> int:
        return self._dim

    def add(self, vector: list[float] | np.ndarray, memory_id: bytes) -> int:
        vec = np.asarray(vector, dtype=np.float32)
        if vec.shape[0] != self._dim:
            raise ValueError(f"Expected dim={self._dim}, got {vec.shape[0]}")

        with self._lock:
            if self.size >= self._max_elements:
                raise ValueError("HNSW index capacity reached")

            if memory_id in self._id_to_node:
                node_id = self._id_to_node[memory_id]
                self._vectors[node_id] = vec
                return node_id

            node_id = len(self._vectors)
            self._vectors.append(vec)
            self._ids.append(memory_id)
            self._id_to_node[memory_id] = node_id

            level = self._random_level()
            while len(self._graphs) <= level:
                self._graphs.append({})

            if self._entry_point < 0:
                self._entry_point = node_id
                self._max_level = level
                for layer in range(level + 1):
                    self._graphs[layer][node_id] = []
                return node_id

            entry_point = self._entry_point
            for layer in range(self._max_level, level, -1):
                entry_point = self._search_layer_greedy(vec, entry_point, layer)

            for layer in range(min(level, self._max_level), -1, -1):
                candidates = self._search_layer(vec, entry_point, self._ef_construction, layer)
                max_conn = self._M if layer > 0 else self._M_max0
                neighbors = self._select_neighbors(vec, candidates, max_conn)
                self._graphs[layer][node_id] = neighbors

                for neighbor in neighbors:
                    nlist = list(self._graphs[layer].get(neighbor, []))
                    nlist.append(node_id)
                    if len(nlist) > max_conn:
                        nlist = self._select_neighbors(self._vectors[neighbor], nlist, max_conn)
                    self._graphs[layer][neighbor] = nlist

                if candidates:
                    entry_point = candidates[0][1]

            if level > self._max_level:
                self._max_level = level
                self._entry_point = node_id

            return node_id

    def search(self, query: list[float] | np.ndarray, k: int = 10) -> list[tuple[bytes, float]]:
        if self._entry_point < 0 or k <= 0:
            return []

        qvec = np.asarray(query, dtype=np.float32)
        if qvec.shape[0] != self._dim:
            return []

        with self._lock:
            entry_point = self._entry_point
            for layer in range(self._max_level, 0, -1):
                entry_point = self._search_layer_greedy(qvec, entry_point, layer)

            candidates = self._search_layer(qvec, entry_point, max(self._ef_search, k), 0)

        results: list[tuple[bytes, float]] = []
        for dist, node_id in candidates[:k]:
            similarity = max(0.0, min(1.0, 1.0 - float(dist)))
            results.append((self._ids[node_id], similarity))
        return results

    def _random_level(self) -> int:
        level = 0
        while random.random() < (1.0 / self._M) and level < 16:
            level += 1
        return level

    @staticmethod
    def _distance(a: np.ndarray, b: np.ndarray) -> float:
        dot = float(np.dot(a, b))
        na = float(np.linalg.norm(a))
        nb = float(np.linalg.norm(b))
        if na < 1e-8 or nb < 1e-8:
            return 1.0
        return 1.0 - (dot / (na * nb))

    def _search_layer_greedy(self, query: np.ndarray, ep: int, layer: int) -> int:
        current = ep
        current_dist = self._distance(query, self._vectors[current])

        while True:
            changed = False
            for neighbor in self._graphs[layer].get(current, []):
                dist = self._distance(query, self._vectors[neighbor])
                if dist < current_dist:
                    current = neighbor
                    current_dist = dist
                    changed = True
            if not changed:
                break
        return current

    def _search_layer(self, query: np.ndarray, ep: int, ef: int, layer: int) -> list[tuple[float, int]]:
        visited = {ep}
        candidates = [(self._distance(query, self._vectors[ep]), ep)]
        best = list(candidates)
        heapq.heapify(candidates)

        while candidates:
            dist_c, candidate = heapq.heappop(candidates)
            worst_best = best[-1][0] if best else float("inf")
            if len(best) >= ef and dist_c > worst_best:
                break

            for neighbor in self._graphs[layer].get(candidate, []):
                if neighbor in visited:
                    continue
                visited.add(neighbor)

                dist = self._distance(query, self._vectors[neighbor])
                if len(best) < ef or dist < best[-1][0]:
                    heapq.heappush(candidates, (dist, neighbor))
                    best.append((dist, neighbor))
                    best.sort(key=lambda item: item[0])
                    if len(best) > ef:
                        best.pop()

        best.sort(key=lambda item: item[0])
        return best

    def _select_neighbors(
        self,
        query: np.ndarray,
        candidates: list[tuple[float, int]] | list[int],
        max_neighbors: int,
    ) -> list[int]:
        if not candidates:
            return []

        if isinstance(candidates[0], tuple):
            scored = list(candidates)  # type: ignore[arg-type]
        else:
            scored = [(self._distance(query, self._vectors[node]), node) for node in candidates]  # type: ignore[misc]

        scored.sort(key=lambda item: item[0])
        return [node for _, node in scored[: max(1, max_neighbors)]]

    def save(self, path: str | Path) -> None:
        vectors = np.vstack(self._vectors).astype(np.float32) if self._vectors else np.empty((0, self._dim), dtype=np.float32)
        metadata = {
            "schema_version": 1,
            "dim": self._dim,
            "M": self._M,
            "ef_construction": self._ef_construction,
            "ef_search": self._ef_search,
            "max_elements": self._max_elements,
            "ids": [base64.b64encode(mid).decode("ascii") for mid in self._ids],
            "graphs": self._graphs,
            "entry_point": self._entry_point,
            "max_level": self._max_level,
        }
        with open(path, "wb") as handle:
            np.savez_compressed(handle, vectors=vectors, metadata=np.array(json.dumps(metadata), dtype=np.str_))

    @classmethod
    def load(cls, path: str | Path) -> "HNSWIndex":
        with np.load(path, allow_pickle=False) as payload:
            metadata_raw = payload["metadata"]
            if isinstance(metadata_raw, np.ndarray):
                metadata_text = str(metadata_raw.item())
            else:
                metadata_text = str(metadata_raw)
            metadata = json.loads(metadata_text)
            vectors = np.asarray(payload["vectors"], dtype=np.float32)

        if not isinstance(metadata, dict):
            raise ValueError("Invalid HNSW metadata payload")

        idx = cls(
            dim=int(metadata["dim"]),
            M=int(metadata["M"]),
            ef_construction=int(metadata.get("ef_construction", 200)),
            ef_search=int(metadata.get("ef_search", 50)),
            max_elements=int(metadata.get("max_elements", 1_000_000)),
        )

        if vectors.ndim != 2 or vectors.shape[1] != idx._dim:
            raise ValueError("Invalid HNSW vector matrix shape")
        idx._vectors = [vectors[i].copy() for i in range(vectors.shape[0])]

        ids_raw = metadata.get("ids", [])
        if not isinstance(ids_raw, list):
            raise ValueError("Invalid HNSW ids payload")
        idx._ids = [base64.b64decode(str(mid)) for mid in ids_raw]
        if len(idx._ids) != len(idx._vectors):
            raise ValueError("HNSW ids/vector length mismatch")
        idx._id_to_node = {mid: i for i, mid in enumerate(idx._ids)}

        raw_graphs = metadata.get("graphs", [])
        if not isinstance(raw_graphs, list):
            raise ValueError("Invalid HNSW graph payload")
        graphs: list[dict[int, list[int]]] = []
        for raw_layer in raw_graphs:
            if not isinstance(raw_layer, dict):
                raise ValueError("Invalid HNSW graph layer payload")
            layer: dict[int, list[int]] = {}
            for node_id, neighbors in raw_layer.items():
                if not isinstance(neighbors, list):
                    raise ValueError("Invalid HNSW graph neighbors payload")
                layer[int(node_id)] = [int(neighbor) for neighbor in neighbors]
            graphs.append(layer)
        idx._graphs = graphs
        idx._entry_point = int(metadata["entry_point"])
        idx._max_level = int(metadata["max_level"])
        return idx
