from pathlib import Path
from typing import List, Dict, Any, Tuple
import json
import faiss
import numpy as np
import pickle


class VectorStore:
    """
    Lightweight FAISS wrapper for ChunkHive.

    Responsibilities:
    - Load FAISS index
    - Load associated chunk metadata
    - Perform similarity search
    """

    def __init__(self, index_path: Path):
        self.index_path = Path(index_path)
        self.index = None
        self.chunks: List[Dict[str, Any]] = []
        self.dimension: int = 0

    # ---------------------------------------------------------
    # Loading
    # ---------------------------------------------------------

    def load(self) -> None:
        """Load FAISS index and chunk metadata."""
        faiss_file = self.index_path.with_suffix(".faiss")
        chunks_file = self.index_path.with_suffix(".chunks")

        if not faiss_file.exists():
            raise FileNotFoundError(f"FAISS index not found: {faiss_file}")

        if not chunks_file.exists():
            raise FileNotFoundError(f"Chunk metadata not found: {chunks_file}")

        # Load FAISS index
        self.index = faiss.read_index(str(faiss_file))
        self.dimension = self.index.d

        # Load chunk metadata
        # with open(chunks_file, "r", encoding="utf-8") as f:
        #     self.chunks = json.load(f)
        with open(chunks_file, "rb") as f:
            self.chunks = pickle.load(f)

        if len(self.chunks) != self.index.ntotal:
            raise ValueError(
                "Mismatch between FAISS vectors and chunk metadata."
            )

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

    def search(
        self,
        query_vector: np.ndarray,
        top_k: int = 5
    ) -> List[Tuple[Dict[str, Any], float]]:
        """
        Perform similarity search.

        Returns:
            List of (chunk, distance)
        """
        if self.index is None:
            raise RuntimeError("VectorStore not loaded. Call load() first.")

        if query_vector.ndim == 1:
            query_vector = query_vector.reshape(1, -1)

        if query_vector.shape[1] != self.dimension:
            raise ValueError(
                f"Query vector dimension {query_vector.shape[1]} "
                f"does not match index dimension {self.dimension}"
            )

        distances, indices = self.index.search(query_vector, top_k)

        results = []
        for idx, dist in zip(indices[0], distances[0]):
            if idx == -1:
                continue
            results.append((self.chunks[idx], float(dist)))

        return results

    # ---------------------------------------------------------
    # Info
    # ---------------------------------------------------------

    def info(self) -> Dict[str, Any]:
        """Return index statistics."""
        if self.index is None:
            raise RuntimeError("VectorStore not loaded.")

        return {
            "total_vectors": self.index.ntotal,
            "dimension": self.dimension,
        }
