"""
H-MEM Level 0: SQLite-backed Episode Store.

Provides persistent storage for agent episodes with efficient
querying by session, agent, and consolidation status.
"""

import json
import sqlite3
from pathlib import Path
from typing import List, Optional

from gsd_rlm.memory.hmem.episode import Episode, EpisodeType


class EpisodeStore:
    """
    SQLite-backed persistent storage for episodes (MEM-01, MEM-11).

    Provides CRUD operations for episodes with indexing on:
    - session_id: for retrieving session history
    - agent_id: for agent-specific memory retrieval
    - timestamp: for time-based queries
    - trace_id: for consolidation status (NULL = unconsolidated)
    """

    def __init__(self, db_path: str = "memory/hmem.db"):
        """
        Initialize the episode store.

        Args:
            db_path: Path to SQLite database file. Parent directories
                     will be created if they don't exist.
        """
        self.db_path = Path(db_path)
        self._init_db()

    def _get_connection(self) -> sqlite3.Connection:
        """
        Get a new database connection.

        Returns:
            SQLite connection with row factory enabled.
        """
        return sqlite3.connect(str(self.db_path))

    def _init_db(self) -> None:
        """Initialize database schema with episodes table and indexes."""
        # Ensure parent directory exists
        self.db_path.parent.mkdir(parents=True, exist_ok=True)

        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # Create episodes table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS episodes (
                    episode_id TEXT PRIMARY KEY,
                    agent_id TEXT NOT NULL,
                    session_id TEXT NOT NULL,
                    episode_type TEXT NOT NULL,
                    context TEXT DEFAULT '',
                    action TEXT DEFAULT '',
                    outcome TEXT DEFAULT '',
                    success INTEGER DEFAULT 0,
                    timestamp TEXT NOT NULL,
                    tokens_used INTEGER DEFAULT 0,
                    duration_ms INTEGER DEFAULT 0,
                    embedding BLOB,
                    tags TEXT DEFAULT '[]',
                    trace_id TEXT
                )
            """)

            # Create indexes for efficient queries
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_episodes_session_id
                ON episodes(session_id)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_episodes_agent_id
                ON episodes(agent_id)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_episodes_timestamp
                ON episodes(timestamp)
            """)
            cursor.execute("""
                CREATE INDEX IF NOT EXISTS idx_episodes_trace_id
                ON episodes(trace_id)
            """)

            conn.commit()
        finally:
            conn.close()

    def store(self, episode: Episode) -> None:
        """
        Store an episode in the database (INSERT OR REPLACE).

        Args:
            episode: Episode to persist.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()

            # Serialize embedding and tags to JSON
            embedding_blob = (
                json.dumps(episode.embedding).encode("utf-8")
                if episode.embedding
                else None
            )
            tags_json = json.dumps(episode.tags)

            cursor.execute(
                """
                INSERT OR REPLACE INTO episodes (
                    episode_id, agent_id, session_id, episode_type,
                    context, action, outcome, success,
                    timestamp, tokens_used, duration_ms,
                    embedding, tags, trace_id
                ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                """,
                (
                    episode.episode_id,
                    episode.agent_id,
                    episode.session_id,
                    episode.episode_type.value,
                    episode.context,
                    episode.action,
                    episode.outcome,
                    1 if episode.success else 0,
                    episode.timestamp,
                    episode.tokens_used,
                    episode.duration_ms,
                    embedding_blob,
                    tags_json,
                    episode.trace_id,
                ),
            )

            conn.commit()
        finally:
            conn.close()

    def get(self, episode_id: str) -> Optional[Episode]:
        """
        Retrieve a single episode by ID.

        Args:
            episode_id: Unique episode identifier.

        Returns:
            Episode if found, None otherwise.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT * FROM episodes WHERE episode_id = ?",
                (episode_id,),
            )
            row = cursor.fetchone()
            if row:
                return self._row_to_episode(row)
            return None
        finally:
            conn.close()

    def get_episodes_for_session(self, session_id: str) -> List[Episode]:
        """
        Get all episodes for a session, ordered by timestamp.

        Args:
            session_id: Session identifier.

        Returns:
            List of episodes for the session, oldest first.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM episodes
                WHERE session_id = ?
                ORDER BY timestamp ASC
                """,
                (session_id,),
            )
            rows = cursor.fetchall()
            return [self._row_to_episode(row) for row in rows]
        finally:
            conn.close()

    def get_episodes_by_agent(self, agent_id: str, limit: int = 50) -> List[Episode]:
        """
        Get recent episodes for an agent.

        Args:
            agent_id: Agent identifier.
            limit: Maximum number of episodes to return.

        Returns:
            List of episodes for the agent, most recent first.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM episodes
                WHERE agent_id = ?
                ORDER BY timestamp DESC
                LIMIT ?
                """,
                (agent_id, limit),
            )
            rows = cursor.fetchall()
            return [self._row_to_episode(row) for row in rows]
        finally:
            conn.close()

    def get_unconsolidated_episodes(self, limit: int = 100) -> List[Episode]:
        """
        Get episodes that haven't been consolidated into higher-level memory.

        Episodes without a trace_id are considered unconsolidated.

        Args:
            limit: Maximum number of episodes to return.

        Returns:
            List of unconsolidated episodes, oldest first.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM episodes
                WHERE trace_id IS NULL
                ORDER BY timestamp ASC
                LIMIT ?
                """,
                (limit,),
            )
            rows = cursor.fetchall()
            return [self._row_to_episode(row) for row in rows]
        finally:
            conn.close()

    def get_episodes_by_time_range(
        self, start_time: str, end_time: str
    ) -> List[Episode]:
        """
        Get episodes within a time range.

        Args:
            start_time: ISO format start timestamp.
            end_time: ISO format end timestamp.

        Returns:
            List of episodes in the time range.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                """
                SELECT * FROM episodes
                WHERE timestamp >= ? AND timestamp <= ?
                ORDER BY timestamp ASC
                """,
                (start_time, end_time),
            )
            rows = cursor.fetchall()
            return [self._row_to_episode(row) for row in rows]
        finally:
            conn.close()

    def delete(self, episode_id: str) -> bool:
        """
        Delete an episode from the store.

        Args:
            episode_id: Episode to delete.

        Returns:
            True if episode was deleted, False if not found.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                "DELETE FROM episodes WHERE episode_id = ?",
                (episode_id,),
            )
            deleted = cursor.rowcount > 0
            conn.commit()
            return deleted
        finally:
            conn.close()

    def count(self) -> int:
        """
        Get total number of episodes in the store.

        Returns:
            Total episode count.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute("SELECT COUNT(*) FROM episodes")
            return cursor.fetchone()[0]
        finally:
            conn.close()

    def count_for_session(self, session_id: str) -> int:
        """
        Count episodes for a specific session.

        Args:
            session_id: Session identifier.

        Returns:
            Number of episodes in the session.
        """
        conn = self._get_connection()
        try:
            cursor = conn.cursor()
            cursor.execute(
                "SELECT COUNT(*) FROM episodes WHERE session_id = ?",
                (session_id,),
            )
            return cursor.fetchone()[0]
        finally:
            conn.close()

    def _row_to_episode(self, row: tuple) -> Episode:
        """
        Convert a database row to an Episode dataclass.

        Args:
            row: SQLite row tuple.

        Returns:
            Episode instance.
        """
        (
            episode_id,
            agent_id,
            session_id,
            episode_type,
            context,
            action,
            outcome,
            success,
            timestamp,
            tokens_used,
            duration_ms,
            embedding_blob,
            tags_json,
            trace_id,
        ) = row

        # Deserialize embedding and tags
        embedding = None
        if embedding_blob:
            try:
                embedding = json.loads(embedding_blob.decode("utf-8"))
            except (json.JSONDecodeError, UnicodeDecodeError):
                pass

        tags = []
        if tags_json:
            try:
                tags = json.loads(tags_json)
            except json.JSONDecodeError:
                pass

        return Episode(
            episode_id=episode_id,
            agent_id=agent_id,
            session_id=session_id,
            episode_type=EpisodeType(episode_type),
            context=context or "",
            action=action or "",
            outcome=outcome or "",
            success=bool(success),
            timestamp=timestamp,
            tokens_used=tokens_used or 0,
            duration_ms=duration_ms or 0,
            embedding=embedding,
            tags=tags,
            trace_id=trace_id,
        )
