"""
PharmaAgent — Clinical Decision RL Environment
models.py: Pydantic types for Action, Observation, State
"""

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


# ── ACTION ────────────────────────────────────────────────────
class Action(BaseModel):
    """
    The agent's action at each step of the clinical decision pipeline.

    action_type options:
      - "diagnose"   : agent proposes a diagnosis from symptoms
      - "select_drug": agent selects a drug to add to the regimen
      - "check_ddi"  : agent checks interaction between two drugs
      - "finalize"   : agent finalises the treatment regimen
    """
    action_type: str = Field(..., description="One of: diagnose, select_drug, check_ddi, finalize")
    value: str = Field(..., description="The agent's response value for the chosen action type")


# ── OBSERVATION ───────────────────────────────────────────────
class Observation(BaseModel):
    """What the agent sees after each step."""
    step: int = Field(..., description="Current step number (1-indexed)")
    phase: str = Field(..., description="Current phase: triage | selection | safety | finalize")
    patient_case: Dict[str, Any] = Field(..., description="The patient case details")
    feedback: str = Field(..., description="Feedback from the environment on the last action")
    valid_options: List[str] = Field(default_factory=list, description="Suggested valid options for the agent")
    reward_so_far: float = Field(default=0.0, description="Cumulative reward accumulated so far")
    done: bool = Field(default=False, description="Whether the episode is complete")


# ── STATE ─────────────────────────────────────────────────────
class State(BaseModel):
    """Internal environment state (not directly exposed to agent)."""
    patient_case: Dict[str, Any] = Field(..., description="Full patient case including ground truth")
    current_phase: str = Field(default="triage", description="triage | selection | safety | finalize")
    step: int = Field(default=0)
    proposed_diagnosis: Optional[str] = Field(default=None)
    selected_drugs: List[str] = Field(default_factory=list)
    checked_interactions: List[Dict[str, Any]] = Field(default_factory=list)
    cumulative_reward: float = Field(default=0.0)
    done: bool = Field(default=False)
    phase_rewards: Dict[str, float] = Field(default_factory=dict)
    max_steps: int = Field(default=8)