chopratejas--headroom/headroom/memory/adapters/sqlite.py

"""SQLite memory store for Headroom's hierarchical memory system.

Provides persistent storage for Memory objects with full support for:
- Hierarchical scope filtering (user/session/agent/turn)
- Temporal versioning with supersession chains
- Point-in-time queries
- Efficient batch operations
"""

from __future__ import annotations

import json
import re
import sqlite3
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any

from ..models import Memory, ScopeLevel
from ..ports import MemoryFilter

if TYPE_CHECKING:
    import numpy as np

# Regex pattern for safe metadata keys: alphanumeric, underscores, hyphens only
# This prevents JSON path injection attacks via malicious key names
_SAFE_METADATA_KEY_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_\-]*$")


def _validate_metadata_key(key: str) -> bool:
    """Validate that a metadata key is safe for use in JSON path expressions.

    Prevents JSON path injection by ensuring keys contain only safe characters.
    Valid keys: start with letter or underscore, contain only alphanumeric, underscore, hyphen.

    Args:
        key: The metadata key to validate.

    Returns:
        True if the key is safe, False otherwise.
    """
    if not key or len(key) > 255:
        return False
    return _SAFE_METADATA_KEY_PATTERN.match(key) is not None


class SQLiteMemoryStore:
    """SQLite-based memory store implementing the MemoryStore protocol.

    Features:
    - Full CRUD operations with batch support
    - Hierarchical scope filtering (user -> session -> agent -> turn)
    - Temporal versioning with supersession chains
    - Point-in-time queries via valid_at filter
    - Thread-safe: connection-per-request pattern

    Usage:
        store = SQLiteMemoryStore("./memories.db")
        await store.save(memory)
        memories = await store.query(MemoryFilter(user_id="alice"))

    Schema:
        The memories table stores all Memory fields with appropriate indexes
        for efficient querying by scope, category, importance, and time.
    """

    def __init__(self, db_path: str | Path = "headroom_memory.db") -> None:
        """Initialize the SQLite memory store.

        Args:
            db_path: Path to SQLite database file. Created if it doesn't exist.
        """
        self.db_path = Path(db_path)
        self._init_db()

    def _get_conn(self) -> sqlite3.Connection:
        """Get a new database connection (thread-safe pattern).

        Returns:
            A new SQLite connection with row factory configured.
        """
        conn = sqlite3.connect(str(self.db_path))
        conn.row_factory = sqlite3.Row
        return conn

    def _init_db(self) -> None:
        """Initialize the database schema with indexes."""
        with self._get_conn() as conn:
            # Create memories table
            conn.execute("""
                CREATE TABLE IF NOT EXISTS memories (
                    id TEXT PRIMARY KEY,
                    content TEXT NOT NULL,

                    -- Hierarchical scoping
                    user_id TEXT NOT NULL,
                    session_id TEXT,
                    agent_id TEXT,
                    turn_id TEXT,

                    -- Temporal
                    created_at TEXT NOT NULL,
                    valid_from TEXT NOT NULL,
                    valid_until TEXT,

                    -- Classification
                    category TEXT NOT NULL,
                    importance REAL NOT NULL DEFAULT 0.5,

                    -- Lineage
                    supersedes TEXT,
                    superseded_by TEXT,
                    promoted_from TEXT,
                    promotion_chain TEXT NOT NULL DEFAULT '[]',

                    -- Access tracking
                    access_count INTEGER NOT NULL DEFAULT 0,
                    last_accessed TEXT,

                    -- Entity references (JSON array)
                    entity_refs TEXT NOT NULL DEFAULT '[]',

                    -- Embedding (BLOB for numpy array)
                    embedding BLOB,

                    -- Metadata (JSON object)
                    metadata TEXT NOT NULL DEFAULT '{}'
                )
            """)

            # Create indexes for efficient querying
            conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_user_id ON memories(user_id)")
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_memories_session_id ON memories(session_id)"
            )
            conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_agent_id ON memories(agent_id)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_turn_id ON memories(turn_id)")
            conn.execute("CREATE INDEX IF NOT EXISTS idx_memories_category ON memories(category)")
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_memories_importance ON memories(importance)"
            )
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_memories_created_at ON memories(created_at)"
            )
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_memories_valid_until ON memories(valid_until)"
            )

            # Composite index for common scope queries
            conn.execute("""
                CREATE INDEX IF NOT EXISTS idx_memories_scope
                ON memories(user_id, session_id, agent_id, turn_id)
            """)

            # Index for supersession chain traversal
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_memories_supersedes ON memories(supersedes)"
            )
            conn.execute(
                "CREATE INDEX IF NOT EXISTS idx_memories_superseded_by ON memories(superseded_by)"
            )

            conn.commit()

    def _serialize_embedding(self, embedding: np.ndarray | None) -> bytes | None:
        """Serialize numpy array to bytes for BLOB storage."""
        if embedding is None:
            return None
        import numpy as np

        return bytes(embedding.astype(np.float32).tobytes())

    def _deserialize_embedding(
        self, data: bytes | None, dim: int | None = None
    ) -> np.ndarray | None:
        """Deserialize bytes back to numpy array."""
        if data is None:
            return None
        import numpy as np

        arr = np.frombuffer(data, dtype=np.float32)
        return arr

    def _memory_to_row(self, memory: Memory) -> dict[str, Any]:
        """Convert Memory object to row dict for insertion."""
        return {
            "id": memory.id,
            "content": memory.content,
            "user_id": memory.user_id,
            "session_id": memory.session_id,
            "agent_id": memory.agent_id,
            "turn_id": memory.turn_id,
            "created_at": memory.created_at.isoformat(),
            "valid_from": memory.valid_from.isoformat(),
            "valid_until": memory.valid_until.isoformat() if memory.valid_until else None,
            "category": "",  # Deprecated - kept for backwards compatibility
            "importance": memory.importance,
            "supersedes": memory.supersedes,
            "superseded_by": memory.superseded_by,
            "promoted_from": memory.promoted_from,
            "promotion_chain": json.dumps(memory.promotion_chain),
            "access_count": memory.access_count,
            "last_accessed": memory.last_accessed.isoformat() if memory.last_accessed else None,
            "entity_refs": json.dumps(memory.entity_refs),
            "embedding": self._serialize_embedding(memory.embedding),
            "metadata": json.dumps(memory.metadata),
        }

    def _row_to_memory(self, row: sqlite3.Row) -> Memory:
        """Convert database row to Memory object."""
        return Memory(
            id=row["id"],
            content=row["content"],
            user_id=row["user_id"],
            session_id=row["session_id"],
            agent_id=row["agent_id"],
            turn_id=row["turn_id"],
            created_at=datetime.fromisoformat(row["created_at"]),
            valid_from=datetime.fromisoformat(row["valid_from"]),
            valid_until=datetime.fromisoformat(row["valid_until"]) if row["valid_until"] else None,
            importance=row["importance"],
            supersedes=row["supersedes"],
            superseded_by=row["superseded_by"],
            promoted_from=row["promoted_from"],
            promotion_chain=json.loads(row["promotion_chain"]) if row["promotion_chain"] else [],
            access_count=row["access_count"],
            last_accessed=datetime.fromisoformat(row["last_accessed"])
            if row["last_accessed"]
            else None,
            entity_refs=json.loads(row["entity_refs"]) if row["entity_refs"] else [],
            embedding=self._deserialize_embedding(row["embedding"]),
            metadata=json.loads(row["metadata"]) if row["metadata"] else {},
        )

    async def save(self, memory: Memory) -> None:
        """Save or update a memory.

        If a memory with the same ID exists, it will be updated.

        Args:
            memory: The memory to save.
        """
        row = self._memory_to_row(memory)

        with self._get_conn() as conn:
            conn.execute(
                """
                INSERT OR REPLACE INTO memories (
                    id, content, user_id, session_id, agent_id, turn_id,
                    created_at, valid_from, valid_until,
                    category, importance,
                    supersedes, superseded_by, promoted_from, promotion_chain,
                    access_count, last_accessed,
                    entity_refs, embedding, metadata
                ) VALUES (
                    :id, :content, :user_id, :session_id, :agent_id, :turn_id,
                    :created_at, :valid_from, :valid_until,
                    :category, :importance,
                    :supersedes, :superseded_by, :promoted_from, :promotion_chain,
                    :access_count, :last_accessed,
                    :entity_refs, :embedding, :metadata
                )
                """,
                row,
            )
            conn.commit()

    async def save_batch(self, memories: list[Memory]) -> None:
        """Save multiple memories in a single transaction.

        Args:
            memories: List of memories to save.
        """
        if not memories:
            return

        rows = [self._memory_to_row(m) for m in memories]

        with self._get_conn() as conn:
            conn.executemany(
                """
                INSERT OR REPLACE INTO memories (
                    id, content, user_id, session_id, agent_id, turn_id,
                    created_at, valid_from, valid_until,
                    category, importance,
                    supersedes, superseded_by, promoted_from, promotion_chain,
                    access_count, last_accessed,
                    entity_refs, embedding, metadata
                ) VALUES (
                    :id, :content, :user_id, :session_id, :agent_id, :turn_id,
                    :created_at, :valid_from, :valid_until,
                    :category, :importance,
                    :supersedes, :superseded_by, :promoted_from, :promotion_chain,
                    :access_count, :last_accessed,
                    :entity_refs, :embedding, :metadata
                )
                """,
                rows,
            )
            conn.commit()

    async def get(self, memory_id: str) -> Memory | None:
        """Retrieve a memory by ID.

        Args:
            memory_id: The unique identifier of the memory.

        Returns:
            The memory if found, None otherwise.
        """
        with self._get_conn() as conn:
            cursor = conn.execute(
                "SELECT * FROM memories WHERE id = ?",
                (memory_id,),
            )
            row = cursor.fetchone()

            if row is None:
                return None

            return self._row_to_memory(row)

    async def get_batch(self, memory_ids: list[str]) -> list[Memory]:
        """Retrieve multiple memories by ID.

        Args:
            memory_ids: List of memory IDs to retrieve.

        Returns:
            List of found memories (may be shorter than input if some not found).
        """
        if not memory_ids:
            return []

        placeholders = ", ".join("?" * len(memory_ids))

        with self._get_conn() as conn:
            cursor = conn.execute(
                f"SELECT * FROM memories WHERE id IN ({placeholders})",  # nosec B608
                memory_ids,
            )

            return [self._row_to_memory(row) for row in cursor]

    async def record_access(
        self,
        memory_ids: list[str],
        accessed_at: datetime | None = None,
    ) -> int:
        """Atomically record one retrieval for each distinct memory ID."""
        unique_ids = list(dict.fromkeys(memory_ids))
        if not unique_ids:
            return 0

        timestamp = accessed_at or datetime.utcnow()
        placeholders = ", ".join("?" for _ in unique_ids)
        with self._get_conn() as conn:
            cursor = conn.execute(
                f"""
                UPDATE memories
                SET access_count = access_count + 1,
                    last_accessed = ?
                WHERE id IN ({placeholders})
                """,  # nosec B608
                [timestamp.isoformat(), *unique_ids],
            )
            conn.commit()
            return cursor.rowcount

    async def delete(self, memory_id: str) -> bool:
        """Delete a memory by ID.

        Args:
            memory_id: The unique identifier of the memory.

        Returns:
            True if the memory was deleted, False if not found.
        """
        with self._get_conn() as conn:
            cursor = conn.execute(
                "DELETE FROM memories WHERE id = ?",
                (memory_id,),
            )
            conn.commit()
            return cursor.rowcount > 0

    async def delete_batch(self, memory_ids: list[str]) -> int:
        """Delete multiple memories by ID.

        Args:
            memory_ids: List of memory IDs to delete.

        Returns:
            Number of memories actually deleted.
        """
        if not memory_ids:
            return 0

        placeholders = ", ".join("?" * len(memory_ids))

        with self._get_conn() as conn:
            cursor = conn.execute(
                f"DELETE FROM memories WHERE id IN ({placeholders})",  # nosec B608
                memory_ids,
            )
            conn.commit()
            return cursor.rowcount

    def _build_query_conditions(self, filter: MemoryFilter) -> tuple[list[str], list[Any]]:
        """Build WHERE clause conditions from a MemoryFilter.

        Returns:
            Tuple of (conditions list, params list).
        """
        conditions: list[str] = []
        params: list[Any] = []

        # Hierarchical scope filtering
        if filter.user_id is not None:
            conditions.append("user_id = ?")
            params.append(filter.user_id)

            # Hierarchical filtering: when filtering by user_id only,
            # return USER-level and below (all that user's memories)
            # This is implicit - we just filter by user_id

            if filter.session_id is not None:
                conditions.append("session_id = ?")
                params.append(filter.session_id)

                if filter.agent_id is not None:
                    conditions.append("agent_id = ?")
                    params.append(filter.agent_id)

                    if filter.turn_id is not None:
                        conditions.append("turn_id = ?")
                        params.append(filter.turn_id)
            elif filter.agent_id is not None:
                # Agent without session - unusual but supported
                conditions.append("agent_id = ?")
                params.append(filter.agent_id)
            elif filter.turn_id is not None:
                # Turn without session/agent - unusual but supported
                conditions.append("turn_id = ?")
                params.append(filter.turn_id)
        elif filter.session_id is not None:
            # Session without user - filter by session only
            conditions.append("session_id = ?")
            params.append(filter.session_id)
        elif filter.agent_id is not None:
            # Agent without user/session
            conditions.append("agent_id = ?")
            params.append(filter.agent_id)
        elif filter.turn_id is not None:
            # Turn only
            conditions.append("turn_id = ?")
            params.append(filter.turn_id)

        # Explicit scope level filtering
        if filter.scope_levels is not None and len(filter.scope_levels) > 0:
            scope_conditions = []
            for level in filter.scope_levels:
                if level == ScopeLevel.USER:
                    # USER level: no session/agent/turn
                    scope_conditions.append(
                        "(session_id IS NULL AND agent_id IS NULL AND turn_id IS NULL)"
                    )
                elif level == ScopeLevel.SESSION:
                    # SESSION level: has session, no agent/turn
                    scope_conditions.append(
                        "(session_id IS NOT NULL AND agent_id IS NULL AND turn_id IS NULL)"
                    )
                elif level == ScopeLevel.AGENT:
                    # AGENT level: has agent, no turn
                    scope_conditions.append("(agent_id IS NOT NULL AND turn_id IS NULL)")
                elif level == ScopeLevel.TURN:
                    # TURN level: has turn
                    scope_conditions.append("(turn_id IS NOT NULL)")

            if scope_conditions:
                conditions.append(f"({' OR '.join(scope_conditions)})")

        # Temporal filtering
        if filter.created_after is not None:
            conditions.append("created_at >= ?")
            params.append(filter.created_after.isoformat())

        if filter.created_before is not None:
            conditions.append("created_at <= ?")
            params.append(filter.created_before.isoformat())

        # Point-in-time query
        if filter.valid_at is not None:
            valid_at_str = filter.valid_at.isoformat()
            conditions.append("valid_from <= ?")
            params.append(valid_at_str)
            conditions.append("(valid_until IS NULL OR valid_until > ?)")
            params.append(valid_at_str)

        # Superseded filtering
        if not filter.include_superseded:
            # Default: only return current memories (not superseded)
            conditions.append("valid_until IS NULL")

        # Importance filtering
        if filter.min_importance is not None:
            conditions.append("importance >= ?")
            params.append(filter.min_importance)

        if filter.max_importance is not None:
            conditions.append("importance <= ?")
            params.append(filter.max_importance)

        # Entity reference filtering (any of the specified entities)
        if filter.entity_refs is not None and len(filter.entity_refs) > 0:
            entity_conditions = []
            for entity_ref in filter.entity_refs:
                # Use JSON contains check
                entity_conditions.append("entity_refs LIKE ?")
                params.append(f'%"{entity_ref}"%')
            conditions.append(f"({' OR '.join(entity_conditions)})")

        # Lineage filtering
        if filter.has_supersedes is not None:
            if filter.has_supersedes:
                conditions.append("supersedes IS NOT NULL")
            else:
                conditions.append("supersedes IS NULL")

        if filter.has_promoted_from is not None:
            if filter.has_promoted_from:
                conditions.append("promoted_from IS NOT NULL")
            else:
                conditions.append("promoted_from IS NULL")

        # Metadata filtering with key validation to prevent JSON path injection
        if filter.metadata_filters:
            for key, value in filter.metadata_filters.items():
                # Validate key to prevent JSON path injection attacks
                # Invalid keys are silently skipped to avoid breaking legitimate queries
                # while blocking malicious attempts like "'] OR 1=1--"
                if not _validate_metadata_key(key):
                    continue
                # Use JSON extraction for metadata filtering
                conditions.append(f"json_extract(metadata, '$.{key}') = ?")
                params.append(json.dumps(value) if not isinstance(value, str) else value)

        return conditions, params

    async def query(self, filter: MemoryFilter) -> list[Memory]:
        """Query memories matching the given filter.

        Args:
            filter: Filter criteria for the query.

        Returns:
            List of matching memories.
        """
        conditions, params = self._build_query_conditions(filter)

        # Build WHERE clause
        where_clause = " AND ".join(conditions) if conditions else "1=1"

        # Build ORDER BY clause
        order_column = filter.order_by
        if order_column not in (
            "created_at",
            "importance",
            "access_count",
            "last_accessed",
        ):
            order_column = "created_at"

        order_direction = "DESC" if filter.order_desc else "ASC"

        # Build full query
        query = f"""
            SELECT * FROM memories
            WHERE {where_clause}
            ORDER BY {order_column} {order_direction}
        """

        # Add pagination using parameterized queries to prevent injection
        if filter.limit is not None:
            query += " LIMIT ?"
            params.append(filter.limit)

        if filter.offset > 0:
            query += " OFFSET ?"
            params.append(filter.offset)

        with self._get_conn() as conn:
            cursor = conn.execute(query, params)
            return [self._row_to_memory(row) for row in cursor]

    async def count(self, filter: MemoryFilter) -> int:
        """Count memories matching the given filter.

        Args:
            filter: Filter criteria for the count.

        Returns:
            Number of matching memories.
        """
        conditions, params = self._build_query_conditions(filter)

        where_clause = " AND ".join(conditions) if conditions else "1=1"

        query = f"SELECT COUNT(*) FROM memories WHERE {where_clause}"

        with self._get_conn() as conn:
            cursor = conn.execute(query, params)
            result = cursor.fetchone()[0]
            return int(result)

    async def supersede(
        self,
        old_memory_id: str,
        new_memory: Memory,
        supersede_time: datetime | None = None,
    ) -> Memory:
        """Supersede an existing memory with a new version.

        This creates a temporal chain: the old memory's valid_until is set,
        and the new memory's supersedes field points to the old one.

        Args:
            old_memory_id: ID of the memory to supersede.
            new_memory: The new memory that replaces it.
            supersede_time: When the supersession occurred (defaults to now).

        Returns:
            The saved new memory with lineage fields populated.

        Raises:
            ValueError: If the old memory is not found.
        """
        if supersede_time is None:
            supersede_time = datetime.now(timezone.utc).replace(tzinfo=None)

        # Get the old memory
        old_memory = await self.get(old_memory_id)
        if old_memory is None:
            raise ValueError(f"Memory with ID {old_memory_id} not found")

        # Update old memory's valid_until and superseded_by
        old_memory.valid_until = supersede_time
        old_memory.superseded_by = new_memory.id

        # Set up new memory's lineage
        new_memory.supersedes = old_memory_id
        new_memory.valid_from = supersede_time

        # Save both in a transaction
        with self._get_conn() as conn:
            # Update old memory
            conn.execute(
                """
                UPDATE memories
                SET valid_until = ?, superseded_by = ?
                WHERE id = ?
                """,
                (supersede_time.isoformat(), new_memory.id, old_memory_id),
            )

            # Insert new memory
            row = self._memory_to_row(new_memory)
            conn.execute(
                """
                INSERT OR REPLACE INTO memories (
                    id, content, user_id, session_id, agent_id, turn_id,
                    created_at, valid_from, valid_until,
                    category, importance,
                    supersedes, superseded_by, promoted_from, promotion_chain,
                    access_count, last_accessed,
                    entity_refs, embedding, metadata
                ) VALUES (
                    :id, :content, :user_id, :session_id, :agent_id, :turn_id,
                    :created_at, :valid_from, :valid_until,
                    :category, :importance,
                    :supersedes, :superseded_by, :promoted_from, :promotion_chain,
                    :access_count, :last_accessed,
                    :entity_refs, :embedding, :metadata
                )
                """,
                row,
            )
            conn.commit()

        return new_memory

    async def get_history(
        self,
        memory_id: str,
        include_future: bool = False,
    ) -> list[Memory]:
        """Get the full history chain for a memory.

        Follows the supersedes/superseded_by chain to return all versions.

        Args:
            memory_id: ID of any memory in the chain.
            include_future: Whether to include memories that superseded this one.

        Returns:
            List of memories in temporal order (oldest first).
        """
        # Start with the given memory
        current = await self.get(memory_id)
        if current is None:
            return []

        history: list[Memory] = [current]

        # Follow chain backwards (supersedes)
        back_id = current.supersedes
        while back_id is not None:
            prev = await self.get(back_id)
            if prev is None:
                break
            history.insert(0, prev)  # Add to beginning
            back_id = prev.supersedes

        # Follow chain forwards (superseded_by) if requested
        if include_future:
            forward_id = current.superseded_by
            while forward_id is not None:
                next_mem = await self.get(forward_id)
                if next_mem is None:
                    break
                history.append(next_mem)
                forward_id = next_mem.superseded_by

        return history

    async def clear_scope(
        self,
        user_id: str,
        session_id: str | None = None,
        agent_id: str | None = None,
        turn_id: str | None = None,
    ) -> int:
        """Clear all memories at or below a scope level.

        Args:
            user_id: Required user scope.
            session_id: If provided, clear session and below.
            agent_id: If provided, clear agent and below.
            turn_id: If provided, clear only that turn.

        Returns:
            Number of memories deleted.
        """
        conditions = ["user_id = ?"]
        params: list[Any] = [user_id]

        if turn_id is not None:
            # Clear only specific turn
            conditions.append("turn_id = ?")
            params.append(turn_id)
        elif agent_id is not None:
            # Clear agent and its turns
            conditions.append("agent_id = ?")
            params.append(agent_id)
        elif session_id is not None:
            # Clear session and its agents/turns
            conditions.append("session_id = ?")
            params.append(session_id)
        # If only user_id, clear all user's memories

        where_clause = " AND ".join(conditions)

        with self._get_conn() as conn:
            cursor = conn.execute(
                f"DELETE FROM memories WHERE {where_clause}",  # nosec B608
                params,
            )
            conn.commit()
            return cursor.rowcount

    async def clear_all(self) -> int:
        """Clear all memories from the store.

        Returns:
            Number of memories deleted.
        """
        with self._get_conn() as conn:
            cursor = conn.execute("DELETE FROM memories")
            conn.commit()
            return cursor.rowcount

    def count_sync(self) -> int:
        """Synchronous count of all memories (for diagnostics).

        Returns:
            Total number of memories in the store.
        """
        with self._get_conn() as conn:
            cursor = conn.execute("SELECT COUNT(*) FROM memories")
            result = cursor.fetchone()[0]
            return int(result)