agents/memory.py

"""
Memory Layer — per-camera embedding ring buffer and event store for ArcisVLM.

Provides temporal context for the agent system:
- EmbeddingRingBuffer: Fixed-size circular buffer storing JEPA embeddings per camera
- EventStore: SQLite-backed persistent event storage with query/aggregation
- MemoryManager: Orchestrates buffers and events across all cameras

The memory layer enables agents to reason about recent scene history,
detect recurring patterns, and retrieve similar past observations via
cosine similarity search over the embedding ring buffer.
"""

from __future__ import annotations

import hashlib
import json
import logging
import sqlite3
import threading
import time
from typing import Any, List, Optional

import torch
import torch.nn.functional as F

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# EmbeddingRingBuffer — fixed-size circular buffer per camera
# ---------------------------------------------------------------------------

class EmbeddingRingBuffer:
    """
    Fixed-size circular buffer that stores the last N JEPA embeddings
    for a single camera stream.

    Supports cosine similarity search for retrieving the most similar
    past embeddings to a query vector.

    Args:
        camera_id: Unique camera identifier
        capacity: Maximum number of embeddings to store
    """

    def __init__(self, camera_id: str, capacity: int = 100) -> None:
        self._camera_id = camera_id
        self._capacity = capacity
        self._buffer: list[dict] = []
        self._write_pos: int = 0
        self._count: int = 0

    @property
    def camera_id(self) -> str:
        return self._camera_id

    @property
    def capacity(self) -> int:
        return self._capacity

    @property
    def size(self) -> int:
        return min(self._count, self._capacity)

    @property
    def is_full(self) -> bool:
        return self._count >= self._capacity

    def push(
        self,
        embedding: torch.Tensor,
        timestamp: float,
        metadata: dict = None,
    ) -> None:
        """
        Add an embedding to the ring buffer.

        When the buffer is full, the oldest entry is overwritten.

        Args:
            embedding: JEPA embedding tensor (any shape, stored detached on CPU)
            timestamp: Wall-clock time of the observation
            metadata: Optional extra data (event type, confidence, etc.)
        """
        entry = {
            "embedding": embedding.detach().cpu(),
            "timestamp": timestamp,
            "metadata": metadata or {},
        }

        if self._count < self._capacity:
            self._buffer.append(entry)
        else:
            self._buffer[self._write_pos] = entry

        self._write_pos = (self._write_pos + 1) % self._capacity
        self._count += 1

    def get_recent(self, n: int = 10) -> List[dict]:
        """
        Return the last N entries in chronological order (oldest first).

        Args:
            n: Number of recent entries to retrieve
        Returns:
            List of dicts with keys: embedding, timestamp, metadata
        """
        actual = min(n, self.size)
        if actual == 0:
            return []

        # Determine indices in reverse insertion order (most recent first)
        results = []
        pos = (self._write_pos - 1) % self._capacity if self.size > 0 else 0
        for _ in range(actual):
            if pos < 0:
                pos = self.size - 1
            results.append(self._buffer[pos])
            pos = (pos - 1) % self.size

        # Return in chronological order (oldest first)
        results.reverse()
        return results

    def similarity_search(
        self,
        query_embedding: torch.Tensor,
        top_k: int = 5,
    ) -> List[dict]:
        """
        Find the top-K most similar embeddings via cosine similarity.

        Args:
            query_embedding: Query embedding tensor
            top_k: Number of results to return
        Returns:
            List of dicts with keys: embedding, timestamp, metadata, similarity
            Sorted by descending similarity.
        """
        if self.size == 0:
            return []

        query = query_embedding.detach().cpu().float()
        if query.dim() > 1:
            query = query.squeeze()

        # Stack all stored embeddings
        stored = torch.stack([
            entry["embedding"].float().squeeze()
            for entry in self._buffer[:self.size]
        ])

        # Cosine similarity: [size]
        similarities = F.cosine_similarity(
            query.unsqueeze(0),
            stored,
            dim=1,
        )

        # Top-K
        k = min(top_k, self.size)
        top_vals, top_indices = similarities.topk(k)

        results = []
        for sim_val, idx in zip(top_vals, top_indices):
            entry = self._buffer[idx.item()]
            results.append({
                "embedding": entry["embedding"],
                "timestamp": entry["timestamp"],
                "metadata": entry["metadata"],
                "similarity": sim_val.item(),
            })

        return results

    def clear(self) -> None:
        """Reset the buffer, discarding all stored embeddings."""
        self._buffer.clear()
        self._write_pos = 0
        self._count = 0


# ---------------------------------------------------------------------------
# EventStore — SQLite-backed event persistence
# ---------------------------------------------------------------------------

class EventStore:
    """
    SQLite-backed persistent event storage for detection events.

    Stores events with camera_id, type, description, confidence,
    optional embedding hash for cross-referencing, and JSON metadata.

    Thread-safe via SQLite's built-in locking + check_same_thread=False.

    Args:
        db_path: Path to SQLite database file, or ":memory:" for in-memory DB
    """

    def __init__(self, db_path: str = "events.db") -> None:
        self.db_path = db_path
        self._conn = sqlite3.connect(db_path, check_same_thread=False)
        self._conn.row_factory = sqlite3.Row
        self._create_tables()

    def _create_tables(self) -> None:
        self._conn.execute("""
            CREATE TABLE IF NOT EXISTS events (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                camera_id TEXT NOT NULL,
                event_type TEXT NOT NULL,
                description TEXT NOT NULL,
                confidence REAL NOT NULL,
                embedding_hash TEXT,
                metadata TEXT,
                timestamp REAL NOT NULL
            )
        """)
        self._conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_events_camera
            ON events (camera_id, timestamp)
        """)
        self._conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_events_type
            ON events (event_type, timestamp)
        """)
        self._conn.commit()

    def record_event(
        self,
        camera_id: str,
        event_type: str,
        description: str,
        confidence: float,
        embedding_hash: str = None,
        metadata: dict = None,
    ) -> int:
        """
        Store a detection event.

        Args:
            camera_id: Source camera
            event_type: Event category (e.g., "person_detected", "vehicle", "anomaly")
            description: Human-readable description
            confidence: Model confidence [0.0, 1.0]
            embedding_hash: Optional hash of the associated embedding
            metadata: Optional JSON-serializable extra data
        Returns:
            Row ID of the inserted event
        """
        cursor = self._conn.execute(
            """INSERT INTO events
               (camera_id, event_type, description, confidence,
                embedding_hash, metadata, timestamp)
               VALUES (?, ?, ?, ?, ?, ?, ?)""",
            (
                camera_id,
                event_type,
                description,
                confidence,
                embedding_hash,
                json.dumps(metadata) if metadata else None,
                time.time(),
            ),
        )
        self._conn.commit()
        return cursor.lastrowid

    def query_events(
        self,
        camera_id: str = None,
        event_type: str = None,
        since: float = None,
        limit: int = 50,
    ) -> List[dict]:
        """
        Query events with optional filters.

        Args:
            camera_id: Filter by camera (None = all cameras)
            event_type: Filter by event type (None = all types)
            since: Only events after this Unix timestamp
            limit: Maximum number of results
        Returns:
            List of event dicts, most recent first
        """
        clauses = []
        params: list[Any] = []

        if camera_id is not None:
            clauses.append("camera_id = ?")
            params.append(camera_id)
        if event_type is not None:
            clauses.append("event_type = ?")
            params.append(event_type)
        if since is not None:
            clauses.append("timestamp >= ?")
            params.append(since)

        where = (" WHERE " + " AND ".join(clauses)) if clauses else ""
        query = f"SELECT * FROM events{where} ORDER BY timestamp DESC LIMIT ?"
        params.append(limit)

        rows = self._conn.execute(query, params).fetchall()
        results = []
        for row in rows:
            d = dict(row)
            if d.get("metadata"):
                d["metadata"] = json.loads(d["metadata"])
            results.append(d)
        return results

    def get_event_counts(
        self,
        camera_id: str = None,
        hours: int = 24,
    ) -> dict:
        """
        Count events by type within a time window.

        Args:
            camera_id: Filter by camera (None = all cameras)
            hours: Look-back window in hours
        Returns:
            Dict mapping event_type -> count
        """
        since = time.time() - (hours * 3600)
        clauses = ["timestamp >= ?"]
        params: list[Any] = [since]

        if camera_id is not None:
            clauses.append("camera_id = ?")
            params.append(camera_id)

        where = " WHERE " + " AND ".join(clauses)
        query = f"SELECT event_type, COUNT(*) as cnt FROM events{where} GROUP BY event_type"

        rows = self._conn.execute(query, params).fetchall()
        return {row["event_type"]: row["cnt"] for row in rows}

    def cleanup_old(self, days: int = 30) -> int:
        """
        Delete events older than the specified number of days.

        Args:
            days: Age threshold in days
        Returns:
            Number of deleted rows
        """
        cutoff = time.time() - (days * 86400)
        cursor = self._conn.execute(
            "DELETE FROM events WHERE timestamp < ?", (cutoff,)
        )
        self._conn.commit()
        return cursor.rowcount

    def close(self) -> None:
        """Close the database connection."""
        self._conn.close()


# ---------------------------------------------------------------------------
# MemoryManager — orchestrates buffers + events
# ---------------------------------------------------------------------------

def _embedding_hash(embedding: torch.Tensor) -> str:
    """Compute a short hash of an embedding for cross-referencing."""
    data = embedding.detach().cpu().numpy().tobytes()
    return hashlib.sha256(data).hexdigest()[:16]


class MemoryManager:
    """
    Central memory orchestrator for the ArcisVLM agent system.

    Lazily creates per-camera EmbeddingRingBuffers and provides a unified
    interface for recording detections and retrieving camera context.

    Args:
        buffer_capacity: Ring buffer size per camera
        db_path: SQLite database path for the event store
    """

    def __init__(
        self,
        buffer_capacity: int = 100,
        db_path: str = "events.db",
    ) -> None:
        self._buffer_capacity = buffer_capacity
        self._buffers: dict[str, EmbeddingRingBuffer] = {}
        self._lock = threading.Lock()
        self.event_store = EventStore(db_path=db_path)

    def get_buffer(self, camera_id: str) -> EmbeddingRingBuffer:
        """
        Get or lazily create the ring buffer for a camera.

        Args:
            camera_id: Camera identifier
        Returns:
            The camera's EmbeddingRingBuffer
        """
        with self._lock:
            if camera_id not in self._buffers:
                self._buffers[camera_id] = EmbeddingRingBuffer(
                    camera_id=camera_id,
                    capacity=self._buffer_capacity,
                )
                logger.info("Created ring buffer for camera '%s'", camera_id)
            return self._buffers[camera_id]

    def record_detection(
        self,
        camera_id: str,
        embedding: torch.Tensor,
        event_type: str,
        description: str,
        confidence: float,
    ) -> None:
        """
        Record a detection: push embedding to ring buffer and log event.

        Args:
            camera_id: Source camera
            embedding: JEPA embedding for this detection
            event_type: Event category
            description: Human-readable description
            confidence: Model confidence [0.0, 1.0]
        """
        timestamp = time.time()
        emb_hash = _embedding_hash(embedding)

        # Push to ring buffer
        buf = self.get_buffer(camera_id)
        buf.push(
            embedding,
            timestamp,
            metadata={
                "event_type": event_type,
                "confidence": confidence,
            },
        )

        # Record in event store
        self.event_store.record_event(
            camera_id=camera_id,
            event_type=event_type,
            description=description,
            confidence=confidence,
            embedding_hash=emb_hash,
            metadata={"embedding_hash": emb_hash},
        )

    def get_scene_descriptor(
        self,
        camera_id: str,
        n_frames: int = 16,
    ) -> Optional[torch.Tensor]:
        """
        Compute a scene descriptor for a camera by averaging recent embeddings.

        This is the primary input to the ConditionEncoder's scene context path.
        Returns the mean of the last N JEPA embeddings from the ring buffer.

        Args:
            camera_id: Camera identifier
            n_frames: Number of recent frames to average (default: 16)
        Returns:
            [embed_dim] tensor — mean embedding, or None if no embeddings stored
        """
        buf = self.get_buffer(camera_id)
        recent = buf.get_recent(n_frames)
        if not recent:
            return None
        embeddings = torch.stack([entry["embedding"].float() for entry in recent])
        if embeddings.dim() > 2:
            embeddings = embeddings.squeeze()
        return embeddings.mean(dim=0)

    def get_camera_context(
        self,
        camera_id: str,
        n_recent: int = 5,
    ) -> dict:
        """
        Build a context snapshot for a camera: recent embeddings + events.

        Useful for providing temporal context to agents making decisions.

        Args:
            camera_id: Camera identifier
            n_recent: Number of recent items to include
        Returns:
            Dict with keys:
                camera_id, recent_embeddings, recent_events, event_counts
        """
        buf = self.get_buffer(camera_id)
        recent_embeddings = buf.get_recent(n_recent)
        recent_events = self.event_store.query_events(
            camera_id=camera_id, limit=n_recent
        )
        event_counts = self.event_store.get_event_counts(camera_id=camera_id)

        return {
            "camera_id": camera_id,
            "recent_embeddings": recent_embeddings,
            "recent_events": recent_events,
            "event_counts": event_counts,
            "buffer_size": buf.size,
            "buffer_full": buf.is_full,
        }