# models.py
# Pydantic models for the study optimizer environment
# Spent way too long deciding whether to use dataclasses or pydantic - went with pydantic
# because the FastAPI integration is just too good to pass up

from pydantic import BaseModel, Field
from typing import Optional


class Observation(BaseModel):
    """What the agent can see about the current state of the student."""

    energy: float = Field(..., ge=0.0, le=100.0, description="Energy level 0-100")
    focus: float = Field(..., ge=0.0, le=100.0, description="Focus score 0-100")
    time_studied: float = Field(..., ge=0.0, description="Total minutes studied this session")
    session_minutes: float = Field(..., ge=0.0, description="Total minutes elapsed in session")
    consecutive_study_minutes: float = Field(
        ..., ge=0.0, description="Minutes studied without a break"
    )
    break_count: int = Field(..., ge=0, description="How many breaks taken so far")
    productive_minutes: float = Field(
        ..., ge=0.0, description="Minutes where focus was above 60"
    )


class Action(BaseModel):
    """The action the agent wants to take."""

    # 0 = study, 1 = short break, 2 = long break
    # took me a while to figure out the right granularity here
    # originally had 5 actions but it made training way slower
    action: int = Field(..., ge=0, le=2, description="0=study, 1=short_break, 2=long_break")


class StepResult(BaseModel):
    """What the environment returns after an action."""

    observation: Observation
    reward: float
    done: bool
    info: dict  # extra stuff for debugging and grading

    # TODO: maybe add a 'truncated' field later for gymnasium compatibility


class TaskInfo(BaseModel):
    """Info about the current task being evaluated."""

    task_id: str
    score: float = Field(..., gt=0.0, lt=1.0)
    done: bool
    extra: Optional[dict] = None