replay/models.py

"""
SQLAlchemy ORM models for the Replay Buffer.

Stores episodes and steps in SQLite for querying, analysis, and training data.
"""

from __future__ import annotations

import json
from datetime import datetime, timezone
from typing import Optional

from sqlalchemy import (
    Boolean, Column, DateTime, Float, ForeignKey, Index,
    Integer, String, Text, create_engine,
)
from sqlalchemy.orm import DeclarativeBase, Session, relationship, sessionmaker


class Base(DeclarativeBase):
    """SQLAlchemy declarative base for all models."""
    pass


class EpisodeRecord(Base):
    """Stores a complete episode of agent-environment interaction.

    Attributes:
        id: Auto-increment primary key.
        episode_id: UUID string for the episode.
        scenario_id: Which scenario was attempted.
        level: Difficulty level (1, 2, or 3).
        total_reward: Sum of all step rewards.
        solved: Whether the scenario was solved.
        total_steps: Number of steps taken.
        timestamp: When the episode occurred.
        training_episode: Which training episode number this was.
    """

    __tablename__ = "episodes"

    id = Column(Integer, primary_key=True, autoincrement=True)
    episode_id = Column(String(36), unique=True, nullable=False, index=True)
    scenario_id = Column(String(100), nullable=False, index=True)
    level = Column(Integer, nullable=False, index=True)
    total_reward = Column(Float, nullable=False, default=0.0)
    solved = Column(Boolean, nullable=False, default=False)
    total_steps = Column(Integer, nullable=False, default=0)
    timestamp = Column(DateTime, nullable=False, default=lambda: datetime.now(timezone.utc))
    training_episode = Column(Integer, nullable=True)

    # Relationship to steps
    steps = relationship("StepRecord", back_populates="episode", cascade="all, delete-orphan",
                          order_by="StepRecord.step_number")

    # Indexes for common queries
    __table_args__ = (
        Index("idx_scenario_solved", "scenario_id", "solved"),
        Index("idx_level_reward", "level", "total_reward"),
        Index("idx_timestamp", "timestamp"),
    )

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "episode_id": self.episode_id,
            "scenario_id": self.scenario_id,
            "level": self.level,
            "total_reward": self.total_reward,
            "solved": self.solved,
            "total_steps": self.total_steps,
            "timestamp": self.timestamp.isoformat() if self.timestamp else None,
            "training_episode": self.training_episode,
            "steps": [s.to_dict() for s in self.steps] if self.steps else [],
        }


class StepRecord(Base):
    """Stores a single step within an episode.

    Attributes:
        id: Auto-increment primary key.
        episode_db_id: Foreign key to EpisodeRecord.
        step_number: Step index within the episode.
        observation_json: JSON-serialized observation dict.
        action: The shell command issued.
        result_json: JSON-serialized ExecutionResult.
        reward: Reward for this step.
        reward_breakdown_json: JSON-serialized reward breakdown.
        error_type: Classified error type at this step.
    """

    __tablename__ = "steps"

    id = Column(Integer, primary_key=True, autoincrement=True)
    episode_db_id = Column(Integer, ForeignKey("episodes.id"), nullable=False)
    step_number = Column(Integer, nullable=False)
    observation_json = Column(Text, nullable=False, default="{}")
    action = Column(Text, nullable=False, default="")
    result_json = Column(Text, nullable=False, default="{}")
    reward = Column(Float, nullable=False, default=0.0)
    reward_breakdown_json = Column(Text, nullable=False, default="{}")
    error_type = Column(String(50), nullable=False, default="unknown")

    episode = relationship("EpisodeRecord", back_populates="steps")

    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "step": self.step_number,
            "observation": json.loads(self.observation_json) if self.observation_json else {},
            "action": self.action,
            "result": json.loads(self.result_json) if self.result_json else {},
            "reward": self.reward,
            "reward_breakdown": json.loads(self.reward_breakdown_json) if self.reward_breakdown_json else {},
            "error_type": self.error_type,
        }


def create_database(db_url: str = "sqlite:///replay_buffer.db") -> sessionmaker:
    """Create the database and return a session factory.

    Args:
        db_url: SQLAlchemy database URL.

    Returns:
        Configured sessionmaker instance.
    """
    engine = create_engine(db_url, echo=False)
    Base.metadata.create_all(engine)
    return sessionmaker(engine, expire_on_commit=False)