"""
ENGRAM Protocol — Manifold Index Tests
Tests for FAISS IndexFlatIP add/search/remove/persist (D2, D4).
"""

from __future__ import annotations

from pathlib import Path

import numpy as np
import pytest
import torch

from kvcos.core.manifold_index import IndexEntry, ManifoldIndex


def _entry(cid: str = "c1", model: str = "llama") -> IndexEntry:
    return IndexEntry(
        cache_id=cid, task_description="test",
        model_id=model, created_at="2026-01-01T00:00:00Z",
        context_len=256, l2_norm=1.0,
    )


class TestAddAndSearch:
    """Add vectors, search via MIPS."""

    def test_add_increments(self) -> None:
        idx = ManifoldIndex(dim=8)
        idx.add(torch.randn(8), _entry("a"))
        idx.add(torch.randn(8), _entry("b"))
        assert idx.n_entries == 2

    def test_search_returns_correct_order(self) -> None:
        idx = ManifoldIndex(dim=4)
        v1 = torch.tensor([1.0, 0.0, 0.0, 0.0])
        v2 = torch.tensor([0.0, 1.0, 0.0, 0.0])
        idx.add(v1, _entry("close"))
        idx.add(v2, _entry("far"))

        query = torch.tensor([1.0, 0.0, 0.0, 0.0])
        results = idx.search(query, top_k=2)
        assert results[0]["cache_id"] == "close"
        assert results[0]["similarity"] > results[1]["similarity"]

    def test_search_empty_returns_empty(self) -> None:
        idx = ManifoldIndex(dim=4)
        results = idx.search(torch.randn(4), top_k=5)
        assert results == []

    def test_model_filter(self) -> None:
        idx = ManifoldIndex(dim=4)
        idx.add(torch.randn(4), _entry("a", model="llama"))
        idx.add(torch.randn(4), _entry("b", model="phi"))
        results = idx.search(torch.randn(4), top_k=10, model_id="phi")
        assert all(r["model_id"] == "phi" for r in results)


class TestRemoveAndRebuild:
    """Remove entries and rebuild index."""

    def test_remove_hides_from_search(self) -> None:
        idx = ManifoldIndex(dim=4)
        v = torch.tensor([1.0, 0.0, 0.0, 0.0])
        idx.add(v, _entry("target"))
        assert idx.remove("target")
        results = idx.search(v, top_k=1)
        assert len(results) == 0

    def test_rebuild_compacts(self) -> None:
        idx = ManifoldIndex(dim=4)
        for i in range(5):
            idx.add(torch.randn(4), _entry(f"c{i}"))
        idx.remove("c1")
        idx.remove("c3")
        active = idx.rebuild()
        assert active == 3


class TestPersistence:
    """Save/load round-trip (D2: serialize_index/deserialize_index)."""

    def test_save_load_round_trip(self, tmp_index_dir: Path) -> None:
        idx = ManifoldIndex(dim=4)
        v1 = torch.tensor([1.0, 0.0, 0.0, 0.0])
        idx.add(v1, _entry("persisted"))
        idx.save(tmp_index_dir / "test.faiss")

        idx2 = ManifoldIndex(dim=4, index_path=tmp_index_dir / "test.faiss")
        assert idx2.n_entries == 1
        results = idx2.search(v1, top_k=1)
        assert results[0]["cache_id"] == "persisted"

    def test_dim_mismatch_raises(self) -> None:
        idx = ManifoldIndex(dim=4)
        with pytest.raises(ValueError, match="dim"):
            idx.add(torch.randn(8), _entry("wrong"))