replay/buffer.py

"""
Replay Buffer — SQLAlchemy-backed episode storage and sampling.

Stores every episode as structured data in SQLite, supports
batch sampling for training, and scenario-based querying for analysis.
"""

from __future__ import annotations

from contextlib import contextmanager
import json
import uuid
import random
from datetime import datetime, timezone
from typing import Dict, Iterator, List, Optional

from sqlalchemy import func, desc
from sqlalchemy.orm import Session

from replay.models import EpisodeRecord, StepRecord, create_database


class ReplayBuffer:
    """Persistent replay buffer backed by SQLite.

    Stores episodes with their steps, supports batch sampling
    for GRPO training, and provides analytics queries.

    Usage:
        buffer = ReplayBuffer("sqlite:///episodes.db")
        buffer.store_episode(episode_data)
        batch = buffer.sample_batch(32)
        scenarios = buffer.get_by_scenario("missing_flask")
    """

    def __init__(self, db_url: str = "sqlite:///replay_buffer.db") -> None:
        """Initialize the replay buffer.

        Args:
            db_url: SQLAlchemy database URL for episode storage.
        """
        self.db_url = db_url
        self._session_factory = create_database(db_url)

    @contextmanager
    def _get_session(self) -> Iterator[Session]:
        """Yield a managed database session and always close it."""
        session = self._session_factory()
        try:
            yield session
        finally:
            session.close()

    def store_episode(
        self,
        scenario_id: str,
        level: int,
        steps: List[Dict],
        total_reward: float,
        solved: bool,
        training_episode: int | None = None,
    ) -> str:
        """Store a complete episode in the buffer.

        Args:
            scenario_id: Which scenario was attempted.
            level: Difficulty level.
            steps: List of step dicts with observation, action, result, reward.
            total_reward: Sum of all step rewards.
            solved: Whether the scenario was solved.
            training_episode: Training episode number, if applicable.

        Returns:
            The generated episode UUID.
        """
        episode_id = str(uuid.uuid4())

        with self._get_session() as session:
            episode = EpisodeRecord(
                episode_id=episode_id,
                scenario_id=scenario_id,
                level=level,
                total_reward=total_reward,
                solved=solved,
                total_steps=len(steps),
                timestamp=datetime.now(timezone.utc),
                training_episode=training_episode,
            )
            session.add(episode)

            for step_data in steps:
                step = StepRecord(
                    episode=episode,
                    step_number=step_data.get("step", 0),
                    observation_json=json.dumps(step_data.get("observation", {})),
                    action=step_data.get("action", ""),
                    result_json=json.dumps(step_data.get("result", {})),
                    reward=step_data.get("reward", 0.0),
                    reward_breakdown_json=json.dumps(step_data.get("reward_breakdown", {})),
                    error_type=step_data.get("error_type", "unknown"),
                )
                session.add(step)

            session.commit()

        return episode_id

    def sample_batch(self, n: int, level: int | None = None) -> List[Dict]:
        """Sample a random batch of episodes for training.

        Args:
            n: Number of episodes to sample.
            level: If provided, only sample from this level.

        Returns:
            List of episode dicts with full step details.
        """
        with self._get_session() as session:
            query = session.query(EpisodeRecord)
            if level is not None:
                query = query.filter(EpisodeRecord.level == level)

            total = query.count()
            if total == 0:
                return []

            # Random offset sampling
            if total <= n:
                episodes = query.all()
            else:
                # Get random IDs
                all_ids = [r.id for r in query.with_entities(EpisodeRecord.id).all()]
                sampled_ids = random.sample(all_ids, min(n, len(all_ids)))
                episodes = query.filter(EpisodeRecord.id.in_(sampled_ids)).all()

            return [ep.to_dict() for ep in episodes]

    def get_by_scenario(self, scenario_id: str, limit: int = 100) -> List[Dict]:
        """Get all episodes for a specific scenario.

        Args:
            scenario_id: Scenario ID to filter by.
            limit: Maximum number of episodes to return.

        Returns:
            List of episode dicts, newest first.
        """
        with self._get_session() as session:
            episodes = (
                session.query(EpisodeRecord)
                .filter(EpisodeRecord.scenario_id == scenario_id)
                .order_by(desc(EpisodeRecord.timestamp))
                .limit(limit)
                .all()
            )
            return [ep.to_dict() for ep in episodes]

    def get_episode(self, episode_id: str) -> Optional[Dict]:
        """Get a specific episode by its UUID.

        Args:
            episode_id: The episode UUID string.

        Returns:
            Episode dict or None if not found.
        """
        with self._get_session() as session:
            episode = (
                session.query(EpisodeRecord)
                .filter(EpisodeRecord.episode_id == episode_id)
                .first()
            )
            if episode:
                return episode.to_dict()
            return None

    def get_stats(self) -> Dict:
        """Get aggregate statistics across all stored episodes.

        Returns:
            Dict with solve rates, mean rewards, counts per level.
        """
        with self._get_session() as session:
            stats: Dict = {"total_episodes": 0, "levels": {}}

            total = session.query(func.count(EpisodeRecord.id)).scalar()
            stats["total_episodes"] = total or 0

            for level in [1, 2, 3]:
                level_query = session.query(EpisodeRecord).filter(EpisodeRecord.level == level)
                level_count = level_query.count()
                if level_count == 0:
                    stats["levels"][level] = {
                        "count": 0, "solve_rate": 0.0,
                        "mean_reward": 0.0, "mean_steps": 0.0,
                    }
                    continue

                solved_count = level_query.filter(EpisodeRecord.solved == True).count()
                mean_reward = session.query(
                    func.avg(EpisodeRecord.total_reward)
                ).filter(EpisodeRecord.level == level).scalar() or 0.0
                mean_steps = session.query(
                    func.avg(EpisodeRecord.total_steps)
                ).filter(EpisodeRecord.level == level).scalar() or 0.0

                stats["levels"][level] = {
                    "count": level_count,
                    "solve_rate": solved_count / level_count if level_count > 0 else 0.0,
                    "mean_reward": round(float(mean_reward), 2),
                    "mean_steps": round(float(mean_steps), 2),
                }

            # Per-scenario stats
            scenario_stats = {}
            scenarios = session.query(
                EpisodeRecord.scenario_id
            ).distinct().all()
            for (sid,) in scenarios:
                sc_query = session.query(EpisodeRecord).filter(EpisodeRecord.scenario_id == sid)
                sc_count = sc_query.count()
                sc_solved = sc_query.filter(EpisodeRecord.solved == True).count()
                scenario_stats[sid] = {
                    "count": sc_count,
                    "solve_rate": sc_solved / sc_count if sc_count > 0 else 0.0,
                }
            stats["scenarios"] = scenario_stats

            return stats

    def get_recent(self, n: int = 20) -> List[Dict]:
        """Get the most recent episodes.

        Args:
            n: Number of episodes to return.

        Returns:
            List of episode dicts, newest first.
        """
        with self._get_session() as session:
            episodes = (
                session.query(EpisodeRecord)
                .order_by(desc(EpisodeRecord.timestamp))
                .limit(n)
                .all()
            )
            return [ep.to_dict() for ep in episodes]

    @property
    def size(self) -> int:
        """Total number of episodes in the buffer."""
        with self._get_session() as session:
            return session.query(func.count(EpisodeRecord.id)).scalar() or 0