from typing import Optional, List
from pydantic import BaseModel, Field


class Action(BaseModel):
    query: str = Field(..., description="The corrected SQL query to execute")
    explanation: Optional[str] = Field(
        default="", description="Optional explanation of what was fixed"
    )


class Observation(BaseModel):
    schema_sql: str = Field(..., description="CREATE TABLE statements for all tables")
    sample_data: str = Field(..., description="INSERT statements showing sample rows")
    buggy_query: str = Field(..., description="The original broken SQL query")
    error_message: str = Field(
        default="",
        description="Error from running the buggy query (empty if semantic bug)",
    )
    task_description: str = Field(..., description="What the query should accomplish")
    expected_columns: List[str] = Field(..., description="Expected result column names")
    attempts_remaining: int = Field(default=5, description="Remaining fix attempts")
    last_attempt_result: str = Field(
        default="", description="Output/error of last submitted query"
    )
    reward: float = Field(default=0.0, description="Reward from last action")
    done: bool = Field(default=False, description="Whether episode is complete")
    success: bool = Field(
        default=False, description="Whether task was solved correctly"
    )


class State(BaseModel):
    episode_id: str = Field(..., description="Unique episode identifier")
    step_count: int = Field(default=0, description="Number of steps taken")
    task_id: str = Field(default="", description="Current task identifier")
    max_steps: int = Field(default=5, description="Maximum steps per episode")
    cumulative_reward: float = Field(default=0.0, description="Total reward so far")