#!/usr/bin/env python3
"""
Local-first vector memory index for Helix Memory Layer V1.

Backed by:
- JSONL for metadata: data/memory_index/meta.jsonl
- NumPy .npy for embeddings: data/memory_index/embeddings.npy

See: ECHO_HELIX_MEMORY_LAYER_V1.md, MEMORY_INDEX_V1.md
"""

import json
from pathlib import Path
from typing import List, Dict, Any, Optional, Union

import numpy as np


class MemoryIndex:
    """
    Local-first vector memory index with cosine similarity search.

    Stores:
    - Metadata as JSONL (one JSON object per line)
    - Embeddings as NumPy .npy array
    """

    def __init__(self, root: Union[Path, str] = Path("data/memory_index")):
        """
        Initialize MemoryIndex.

        Args:
            root: Directory to store index files
        """
        self.root = Path(root)
        self.meta_path = self.root / "meta.jsonl"
        self.emb_path = self.root / "embeddings.npy"

        self._meta: List[Dict[str, Any]] = []
        self._emb: Optional[np.ndarray] = None
        self._dim: Optional[int] = None

        self._load()

    def add(self, item_id: str, embedding: np.ndarray, metadata: Dict[str, Any]) -> None:
        """
        Append a new vector + metadata to the index.

        Args:
            item_id: Unique identifier for this item
            embedding: 1D numpy array (the vector)
            metadata: Arbitrary metadata dict
        """
        # Ensure embedding is numpy array and 1D
        embedding = np.asarray(embedding).flatten()

        # Set dimension on first insert
        if self._dim is None:
            self._dim = embedding.shape[0]
        elif embedding.shape[0] != self._dim:
            raise ValueError(
                f"Embedding dimension mismatch: expected {self._dim}, got {embedding.shape[0]}"
            )

        # Append metadata
        record = {
            "id": item_id,
            "dim": int(self._dim),
            "metadata": metadata,
        }
        self._meta.append(record)

        # Append embedding
        if self._emb is None or self._emb.shape[0] == 0:
            self._emb = embedding.reshape(1, -1)
        else:
            self._emb = np.vstack([self._emb, embedding.reshape(1, -1)])

        # Persist
        self._save()

    def query(self, embedding: np.ndarray, top_k: int = 5) -> List[Dict[str, Any]]:
        """
        Return top_k most similar entries by cosine similarity.

        Args:
            embedding: Query vector (1D numpy array)
            top_k: Number of results to return

        Returns:
            List of dicts with id, score, and metadata
        """
        if self._emb is None or self._emb.shape[0] == 0:
            return []

        # Ensure query is 1D
        q = np.asarray(embedding).flatten()

        # Compute cosine similarity
        # cos_sim = (a · b) / (||a|| * ||b||)
        q_norm = np.linalg.norm(q)
        if q_norm == 0:
            return []

        # Normalize query
        q_normalized = q / q_norm

        # Normalize stored embeddings
        emb_norms = np.linalg.norm(self._emb, axis=1, keepdims=True)
        emb_norms = np.where(emb_norms == 0, 1, emb_norms)  # Avoid division by zero
        emb_normalized = self._emb / emb_norms

        # Dot product gives cosine similarity
        similarities = emb_normalized @ q_normalized

        # Get top-k indices
        if len(similarities) <= top_k:
            top_indices = np.argsort(similarities)[::-1]
        else:
            top_indices = np.argpartition(similarities, -top_k)[-top_k:]
            top_indices = top_indices[np.argsort(similarities[top_indices])[::-1]]

        # Build results
        results = []
        for i in top_indices:
            results.append({
                "id": self._meta[i]["id"],
                "score": float(similarities[i]),
                "metadata": self._meta[i]["metadata"],
            })

        return results

    def get(self, item_id: str) -> Optional[Dict[str, Any]]:
        """
        Get a record by ID.

        Args:
            item_id: The item ID to look up

        Returns:
            Dict with id, embedding, metadata or None if not found
        """
        for i, record in enumerate(self._meta):
            if record["id"] == item_id:
                return {
                    "id": record["id"],
                    "embedding": self._emb[i].tolist() if self._emb is not None else None,
                    "metadata": record["metadata"],
                }
        return None

    def count(self) -> int:
        """Return number of items in index."""
        return len(self._meta)

    def _save(self) -> None:
        """Persist meta + embeddings to disk."""
        # Ensure directory exists
        self.root.mkdir(parents=True, exist_ok=True)

        # Write metadata as JSONL
        with open(self.meta_path, "w") as f:
            for record in self._meta:
                f.write(json.dumps(record, separators=(",", ":")) + "\n")

        # Save embeddings if we have any
        if self._emb is not None and self._emb.shape[0] > 0:
            np.save(self.emb_path, self._emb)

    def _load(self) -> None:
        """Load meta + embeddings from disk if present."""
        # Create directory if missing
        self.root.mkdir(parents=True, exist_ok=True)

        # Load metadata
        if self.meta_path.exists():
            self._meta = []
            with open(self.meta_path, "r") as f:
                for line in f:
                    line = line.strip()
                    if line:
                        self._meta.append(json.loads(line))

            # Infer dimension from first record if available
            if self._meta:
                self._dim = self._meta[0].get("dim")

        # Load embeddings
        if self.emb_path.exists():
            self._emb = np.load(self.emb_path)
            if self._dim is None and self._emb.shape[0] > 0:
                self._dim = self._emb.shape[1]

    def clear(self) -> None:
        """Clear all data from the index."""
        self._meta = []
        self._emb = None
        self._dim = None

        # Remove files if they exist
        if self.meta_path.exists():
            self.meta_path.unlink()
        if self.emb_path.exists():
            self.emb_path.unlink()
