models.py

"""
SQLEnv Pydantic models — the data contracts between client and server.

These models define the typed interface for the SQLEnv RL environment,
following the OpenEnv pattern (see OpenEnv Tutorial for reference):

    Action      — what the agent sends each step
    Observation — what the agent receives back
    State       — episode metadata (exposed via the state endpoint)

RL terminology — state vs observation
─────────────────────────────────────
In RL theory:

    State (s)       A COMPLETE description of the world. Nothing is hidden.
    Observation (o) A PARTIAL description of a state, which may omit info.

In SQLEnv these map to:

    EpisodeContext  The full RL state (s). Lives on the server only.
                    Contains gold answers, reward accumulators, DB
                    connection, full query history — everything needed
                    to advance the simulation and compute rewards.

    SQLObservation  The observation (o). Sent to the agent over the wire.
                    Contains the question, truncated results, revealed
                    schema, budget, and action history. The agent NEVER
                    sees the gold answer, progress scores, or full DB.

    SQLState        OpenEnv's "State" base class — lightweight episode
                    metadata (episode_id, step_count). This is NOT the
                    RL state; it is a convenience for logging/debugging.

This separation is what makes SQLEnv a POMDP: the agent must act under
uncertainty, which is what makes exploration necessary and learnable.
"""

import sqlite3
from dataclasses import dataclass, field as dataclass_field

from openenv.core.env_server.interfaces import Message
from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field
import torch

# ---------------------------------------------------------------------------
# Wire types: these cross the HTTP boundary between client and server
# ---------------------------------------------------------------------------


class SQLAction(Action):
    """What the agent sends each step.

    The action space is intentionally small and structured so agents can
    explicitly control the environment loop.
    """

    action_type: str = Field(
        ...,
        description="One of: DESCRIBE, SAMPLE, QUERY, ANSWER",
    )
    argument: str = Field(
        ...,
        description=(
            "Table name (DESCRIBE/SAMPLE), SQL string (QUERY), "
            "or answer value (ANSWER)."
        ),
    )


class SQLObservation(Observation):
    """What the agent receives after each step.

    This is the agent's PARTIAL view of the world. Key design choices:

    - schema_info starts with table names only; columns are revealed
      incrementally as the agent DESCRIBEs tables.
    - result is always a truncated string, never raw data. The agent sees
      what a human analyst would see in a terminal — at most N rows of
      formatted text. This keeps the observation bounded and forces the
      agent to reason about what it sees rather than brute-force scanning.
    - action_history gives the agent memory of its own trajectory without
      the server needing to re-send full results from prior steps.
    """

    # Inherited from Observation: done (bool), reward (float | None)
    question: str = Field(..., description="The NL question to answer")
    schema_info: str = Field(..., description="Known schema information")
    result: str = Field(default="", description="Result of the last action")
    error: str = Field(default="", description="Error message if action failed")
    step_count: int = Field(default=0, description="Current step number")
    budget_remaining: int = Field(default=0, description="Steps remaining")
    action_history: list[str] = Field(
        default_factory=list,
        description="Summary of previous actions",
    )


class SQLState(State):
    """Episode metadata exposed via GET /state.

    This is the minimal public state — enough for logging and debugging,
    but NOT the full internal bookkeeping (see EpisodeContext below).
    """

    # # Inherited from State: episode_id (str | None), step_count (int)
    # game_name: str = Field(
    #     "sql_env", description="Name of the game/environment"
    # )
    history_messages: list[Message] = Field(default_factory=list)
    history_tokens: list[torch.Tensor] = Field(
        default_factory=list
    )  # Same len as messages
    current_action_type: str = Field(
        default="QUERY",
        description="Current action type: DESCRIBE, SAMPLE, QUERY, or ANSWER",
    )


@dataclass
class QuestionRecord:
    """One question from the Spider dataset."""

    question_id: str
    question_text: str
    database_name: str
    gold_sql: str
    gold_answer: str
    answer_type: str
    difficulty: str
    tables_involved: list[str]


@dataclass
class EpisodeContext:
    """Per-episode server-side state (never sent to agent)."""

    episode_id: str
    db_connection: sqlite3.Connection
    question_record: QuestionRecord
    step_count: int = 0
    budget: int = 15
    described_tables: set[str] = dataclass_field(default_factory=set)
    action_log: list[str] = dataclass_field(default_factory=list)
    done: bool = False
    gold_answer: str | None = None
    gold_rows: list[tuple] = dataclass_field(default_factory=list)
    query_hashes: set[str] = dataclass_field(default_factory=set)
    best_progress: float = 0.0
    cumulative_step_reward: float = 0.0
    cumulative_new_info_reward: float = 0.0


# ---------------------------------------------------------------------------
# Conceptual internal state: what the server tracks per episode
# ---------------------------------------------------------------------------
#
# The classes below are a DESIGN OUTLINE, not runnable implementation.
# They describe the information the server needs to maintain during an
# episode so that it can:
#
#   1. Execute actions against the database
#   2. Compute the 3-layer reward signal
#   3. Enforce budget limits and anti-gaming measures
#   4. Build the next observation for the agent
#
# These are SERVER-ONLY — they never cross the HTTP boundary.
# Implementation will follow in server/environment.py during Phase 2.
#
#
# EpisodeContext — Per-episode server state
# ──────────────────────────────────────────
# Conceptual fields:
#
#   episode_id: str
#       Unique identifier for this episode (UUID).
#
#   question_record: QuestionRecord
#       The selected question and its metadata:
#         - question_id, question_text, database_name
#         - gold_sql, gold_answer, answer_type, difficulty
#       Loaded from the question set JSON at reset().
#
#   db_connection: sqlite3.Connection
#       Read-only connection to the episode's SQLite database.
#       Opened at reset(), closed when the episode ends.
#       Enforces: read-only mode, statement timeout (5s), SELECT-only.
#
#   step_count: int
#       Current step number (0 at reset, incremented each step()).
#
#   budget: int
#       Steps remaining. Starts at max_steps (default 15).
#       Decremented on each non-ANSWER action. Episode terminates
#       when budget hits 0 without an ANSWER.
#
#   --- Schema tracking (for observation building) ---
#
#   known_tables: set[str]
#       Table names revealed to the agent. Starts with ALL table names
#       (agent sees table names at reset), but column details are hidden.
#
#   described_tables: dict[str, list[ColumnInfo]]
#       Tables the agent has DESCRIBEd → their column info.
#       Used to build the incrementally-revealed schema_info string.
#
#   --- Reward tracking (Layer 1: Operational) ---
#
#   query_hashes: set[str]
#       Hashes of all SQL queries executed this episode.
#       Used for repeat detection (r_repeat penalty).
#
#   explored_entities: set[str]
#       Set of "table.column" strings the agent has discovered.
#       Used for r_new_info reward. Capped at 0.10 total per episode.
#
#   cumulative_new_info_reward: float
#       Running total of r_new_info awarded. Once this reaches the cap
#       (0.10), no more r_new_info is given.
#
#   --- Reward tracking (Layer 2: Progress) ---
#
#   gold_result: Any
#       The result of running gold_sql on the database, computed once
#       at reset(). This is the reference for progress comparison.
#
#   best_progress: float
#       Best binned progress score achieved so far (one of
#       {0, 0.25, 0.5, 0.75, 1.0}). Reward is given only when
#       a QUERY result IMPROVES over this value.
#
#   --- Reward tracking (aggregates) ---
#
#   cumulative_step_reward: float
#       Running sum of all per-step rewards (Layers 1 + 2).
#       Clamped to [-0.2, +0.5] at episode end.
#
#   --- Action history (for observation) ---
#
#   action_log: list[str]
#       Human-readable summaries of each action taken, e.g.:
#         "DESCRIBE employees → 5 columns"
#         "QUERY: SELECT COUNT(*) FROM orders → 42"
#         "ANSWER: 42 → correct"
#       Sent to the agent in SQLObservation.action_history so it has
#       memory of its own trajectory.
#
#
# QuestionRecord — Metadata for a single question
# ─────────────────────────────────────────────────
# Conceptual fields:
#
#   question_id: str          e.g. "spider_dev_042"
#   question_text: str        The natural language question
#   database_name: str        Which SQLite database to load
#   gold_sql: str             Reference SQL (hidden from agent)
#   gold_answer: str          Expected answer (hidden from agent)
#   answer_type: str          One of: integer, float, string, list, table
#   difficulty: str           One of: easy, medium, hard
#   tables_involved: list[str]  Which tables the gold query touches
#
#
# ColumnInfo — Schema detail for a single column
# ───────────────────────────────────────────────
# Conceptual fields:
#
#   name: str                 Column name
#   dtype: str                SQLite type (TEXT, INTEGER, REAL, etc.)
#   is_primary_key: bool      Whether this is a PK
#   is_foreign_key: bool      Whether this is a FK
#   references: str | None    "table.column" if FK, else None
#