tests/test_memory.py

"""
Tests for the ArcisVLM memory layer: ring buffer, event store, and memory manager.
"""

import time

import pytest
import torch

from agents.memory import EmbeddingRingBuffer, EventStore, MemoryManager


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _rand_embedding(dim: int = 128) -> torch.Tensor:
    """Generate a random unit-norm embedding."""
    e = torch.randn(dim)
    return e / e.norm()


# ---------------------------------------------------------------------------
# EmbeddingRingBuffer tests
# ---------------------------------------------------------------------------

class TestEmbeddingRingBuffer:

    def test_init(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=10)
        assert buf.camera_id == "cam-1"
        assert buf.capacity == 10
        assert buf.size == 0
        assert not buf.is_full

    def test_push_and_size(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=5)
        for i in range(3):
            buf.push(_rand_embedding(), timestamp=float(i))
        assert buf.size == 3
        assert not buf.is_full

    def test_full_flag(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=3)
        for i in range(3):
            buf.push(_rand_embedding(), timestamp=float(i))
        assert buf.is_full

    def test_overflow_wraps(self):
        """After pushing more than capacity, size stays at capacity."""
        buf = EmbeddingRingBuffer("cam-1", capacity=3)
        for i in range(10):
            buf.push(_rand_embedding(), timestamp=float(i))
        assert buf.size == 3
        assert buf.is_full

    def test_get_recent_returns_chronological(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=10)
        for i in range(5):
            buf.push(_rand_embedding(), timestamp=float(i))

        recent = buf.get_recent(3)
        assert len(recent) == 3
        # Should be in chronological order (oldest first)
        assert recent[0]["timestamp"] < recent[1]["timestamp"] < recent[2]["timestamp"]
        # Should be the last 3 pushed
        assert recent[0]["timestamp"] == 2.0
        assert recent[1]["timestamp"] == 3.0
        assert recent[2]["timestamp"] == 4.0

    def test_get_recent_after_overflow(self):
        """get_recent works correctly after the buffer has wrapped around."""
        buf = EmbeddingRingBuffer("cam-1", capacity=3)
        for i in range(7):
            buf.push(_rand_embedding(), timestamp=float(i))

        recent = buf.get_recent(3)
        assert len(recent) == 3
        # Should contain the last 3 entries: 4, 5, 6
        timestamps = [r["timestamp"] for r in recent]
        assert timestamps == [4.0, 5.0, 6.0]

    def test_get_recent_more_than_available(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=10)
        buf.push(_rand_embedding(), timestamp=1.0)
        buf.push(_rand_embedding(), timestamp=2.0)
        recent = buf.get_recent(100)
        assert len(recent) == 2

    def test_get_recent_empty(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=5)
        assert buf.get_recent(5) == []

    def test_metadata_preserved(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=5)
        buf.push(_rand_embedding(), timestamp=1.0, metadata={"label": "person"})
        recent = buf.get_recent(1)
        assert recent[0]["metadata"] == {"label": "person"}

    def test_clear(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=5)
        for i in range(5):
            buf.push(_rand_embedding(), timestamp=float(i))
        assert buf.size == 5
        buf.clear()
        assert buf.size == 0
        assert not buf.is_full
        assert buf.get_recent(5) == []

    def test_similarity_search_basic(self):
        """Similarity search returns the most similar embedding."""
        buf = EmbeddingRingBuffer("cam-1", capacity=10)

        # Push a known embedding
        target = _rand_embedding(128)
        buf.push(target, timestamp=0.0)

        # Push several random ones
        for i in range(1, 8):
            buf.push(_rand_embedding(128), timestamp=float(i))

        results = buf.similarity_search(target, top_k=1)
        assert len(results) == 1
        assert results[0]["similarity"] > 0.99  # should match itself
        assert results[0]["timestamp"] == 0.0

    def test_similarity_search_ranking(self):
        """Top-K results are ranked by descending similarity."""
        buf = EmbeddingRingBuffer("cam-1", capacity=20)

        # Create a target and a very similar embedding
        target = _rand_embedding(64)
        similar = target + 0.05 * torch.randn(64)
        similar = similar / similar.norm()

        buf.push(target, timestamp=0.0, metadata={"tag": "exact"})
        buf.push(similar, timestamp=1.0, metadata={"tag": "similar"})
        for i in range(2, 10):
            buf.push(_rand_embedding(64), timestamp=float(i))

        results = buf.similarity_search(target, top_k=3)
        assert len(results) == 3
        # First result should be the exact match
        assert results[0]["similarity"] >= results[1]["similarity"]
        assert results[1]["similarity"] >= results[2]["similarity"]
        assert results[0]["metadata"]["tag"] == "exact"

    def test_similarity_search_empty(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=5)
        results = buf.similarity_search(_rand_embedding(), top_k=3)
        assert results == []

    def test_similarity_search_top_k_exceeds_size(self):
        buf = EmbeddingRingBuffer("cam-1", capacity=10)
        buf.push(_rand_embedding(), timestamp=1.0)
        buf.push(_rand_embedding(), timestamp=2.0)
        results = buf.similarity_search(_rand_embedding(), top_k=100)
        assert len(results) == 2


# ---------------------------------------------------------------------------
# EventStore tests
# ---------------------------------------------------------------------------

class TestEventStore:

    def _make_store(self) -> EventStore:
        return EventStore(db_path=":memory:")

    def test_record_and_query(self):
        store = self._make_store()
        store.record_event("cam-1", "person", "Person detected", 0.92)
        events = store.query_events(camera_id="cam-1")
        assert len(events) == 1
        assert events[0]["event_type"] == "person"
        assert events[0]["confidence"] == 0.92

    def test_query_by_event_type(self):
        store = self._make_store()
        store.record_event("cam-1", "person", "Person", 0.9)
        store.record_event("cam-1", "vehicle", "Car", 0.8)
        store.record_event("cam-2", "person", "Person", 0.85)

        persons = store.query_events(event_type="person")
        assert len(persons) == 2

        vehicles = store.query_events(event_type="vehicle")
        assert len(vehicles) == 1

    def test_query_by_camera(self):
        store = self._make_store()
        store.record_event("cam-1", "person", "Person", 0.9)
        store.record_event("cam-2", "vehicle", "Car", 0.8)

        cam1 = store.query_events(camera_id="cam-1")
        assert len(cam1) == 1
        assert cam1[0]["camera_id"] == "cam-1"

    def test_query_since(self):
        store = self._make_store()
        now = time.time()
        store.record_event("cam-1", "person", "Person", 0.9)

        # Query with a past timestamp should include the event
        events = store.query_events(since=now - 10)
        assert len(events) == 1

        # Query with a future timestamp should exclude it
        events = store.query_events(since=now + 100)
        assert len(events) == 0

    def test_query_limit(self):
        store = self._make_store()
        for i in range(20):
            store.record_event("cam-1", "person", f"Person {i}", 0.9)
        events = store.query_events(limit=5)
        assert len(events) == 5

    def test_query_order_most_recent_first(self):
        store = self._make_store()
        store.record_event("cam-1", "person", "First", 0.9)
        time.sleep(0.01)  # ensure different timestamp
        store.record_event("cam-1", "person", "Second", 0.8)

        events = store.query_events()
        assert events[0]["description"] == "Second"
        assert events[1]["description"] == "First"

    def test_metadata_stored(self):
        store = self._make_store()
        store.record_event(
            "cam-1", "anomaly", "Fire detected", 0.95,
            metadata={"severity": "high", "zone": "parking"},
        )
        events = store.query_events()
        assert events[0]["metadata"]["severity"] == "high"

    def test_embedding_hash_stored(self):
        store = self._make_store()
        store.record_event(
            "cam-1", "person", "Person", 0.9,
            embedding_hash="abc123",
        )
        events = store.query_events()
        assert events[0]["embedding_hash"] == "abc123"

    def test_get_event_counts(self):
        store = self._make_store()
        store.record_event("cam-1", "person", "Person", 0.9)
        store.record_event("cam-1", "person", "Person", 0.85)
        store.record_event("cam-1", "vehicle", "Car", 0.8)
        store.record_event("cam-2", "person", "Person", 0.9)

        # All cameras
        counts = store.get_event_counts()
        assert counts["person"] == 3
        assert counts["vehicle"] == 1

        # Single camera
        counts_cam1 = store.get_event_counts(camera_id="cam-1")
        assert counts_cam1["person"] == 2
        assert counts_cam1["vehicle"] == 1

    def test_cleanup_old(self):
        store = self._make_store()
        # Insert events, then manually backdate some
        store.record_event("cam-1", "person", "Recent", 0.9)

        # Insert an old event by manipulating SQL directly
        old_ts = time.time() - (60 * 86400)  # 60 days ago
        store._conn.execute(
            """INSERT INTO events
               (camera_id, event_type, description, confidence, timestamp)
               VALUES (?, ?, ?, ?, ?)""",
            ("cam-1", "person", "Old", 0.7, old_ts),
        )
        store._conn.commit()

        assert len(store.query_events()) == 2
        deleted = store.cleanup_old(days=30)
        assert deleted == 1
        remaining = store.query_events()
        assert len(remaining) == 1
        assert remaining[0]["description"] == "Recent"


# ---------------------------------------------------------------------------
# MemoryManager tests
# ---------------------------------------------------------------------------

class TestMemoryManager:

    def _make_manager(self, capacity: int = 10) -> MemoryManager:
        return MemoryManager(buffer_capacity=capacity, db_path=":memory:")

    def test_get_buffer_creates_lazily(self):
        mgr = self._make_manager()
        buf = mgr.get_buffer("cam-1")
        assert isinstance(buf, EmbeddingRingBuffer)
        assert buf.camera_id == "cam-1"
        # Same buffer returned on second call
        assert mgr.get_buffer("cam-1") is buf

    def test_multiple_cameras(self):
        mgr = self._make_manager()
        buf1 = mgr.get_buffer("cam-1")
        buf2 = mgr.get_buffer("cam-2")
        assert buf1 is not buf2
        assert buf1.camera_id == "cam-1"
        assert buf2.camera_id == "cam-2"

    def test_record_detection(self):
        mgr = self._make_manager()
        emb = _rand_embedding(128)
        mgr.record_detection(
            camera_id="cam-1",
            embedding=emb,
            event_type="person",
            description="Person at entrance",
            confidence=0.92,
        )

        # Check ring buffer got the embedding
        buf = mgr.get_buffer("cam-1")
        assert buf.size == 1

        # Check event store got the event
        events = mgr.event_store.query_events(camera_id="cam-1")
        assert len(events) == 1
        assert events[0]["event_type"] == "person"
        assert events[0]["confidence"] == 0.92

    def test_record_detection_multiple_cameras(self):
        mgr = self._make_manager()
        for cam in ["cam-1", "cam-2", "cam-3"]:
            for i in range(3):
                mgr.record_detection(
                    camera_id=cam,
                    embedding=_rand_embedding(64),
                    event_type="person",
                    description=f"Person {i}",
                    confidence=0.8 + i * 0.05,
                )

        # Each camera should have 3 embeddings
        for cam in ["cam-1", "cam-2", "cam-3"]:
            assert mgr.get_buffer(cam).size == 3

        # Total events across all cameras
        all_events = mgr.event_store.query_events(limit=100)
        assert len(all_events) == 9

    def test_get_camera_context(self):
        mgr = self._make_manager()

        # Record some detections
        for i in range(5):
            mgr.record_detection(
                camera_id="cam-1",
                embedding=_rand_embedding(64),
                event_type="person" if i % 2 == 0 else "vehicle",
                description=f"Event {i}",
                confidence=0.9,
            )

        ctx = mgr.get_camera_context("cam-1", n_recent=3)
        assert ctx["camera_id"] == "cam-1"
        assert len(ctx["recent_embeddings"]) == 3
        assert len(ctx["recent_events"]) == 3
        assert "person" in ctx["event_counts"]
        assert "vehicle" in ctx["event_counts"]
        assert ctx["buffer_size"] == 5
        assert not ctx["buffer_full"]

    def test_get_camera_context_empty_camera(self):
        mgr = self._make_manager()
        ctx = mgr.get_camera_context("cam-new")
        assert ctx["camera_id"] == "cam-new"
        assert ctx["recent_embeddings"] == []
        assert ctx["recent_events"] == []
        assert ctx["event_counts"] == {}
        assert ctx["buffer_size"] == 0

    def test_buffer_capacity_respected(self):
        mgr = self._make_manager(capacity=5)
        for i in range(20):
            mgr.record_detection(
                camera_id="cam-1",
                embedding=_rand_embedding(64),
                event_type="person",
                description=f"Person {i}",
                confidence=0.9,
            )
        assert mgr.get_buffer("cam-1").size == 5
        assert mgr.get_buffer("cam-1").is_full