models.py

from enum import Enum
from typing import Any, Optional

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


class GameState(str, Enum):
    TOSS = "toss"
    BATTING = "batting"
    BOWLING = "bowling"
    FINISHED = "finished"


class CricketAction(Action):
    """Agent sends a structured JSON tool call per step."""

    tool: str = Field(
        ...,
        description="Tool to invoke: 'call_toss' | 'choose_bowler' | 'select_batter' | 'set_strategy' | 'plan_shot' | 'play_delivery' | 'set_bowling_strategy' | 'plan_delivery' | 'set_field_setting' | 'bowl_delivery' | 'reflect_after_ball' | 'analyze_situation' | 'set_match_plan' | 'update_match_plan'",
    )
    arguments: dict[str, Any] = Field(
        default_factory=dict,
        description="Tool-specific arguments",
    )


class CricketObservation(Observation):
    """Environment returns game state + rendered prompt for LLM consumption."""

    game_state: str = Field(default="batting")
    game_context: dict[str, Any] = Field(
        default_factory=dict,
        description="Current match state: over, ball, score, wickets, phase, run_rate, etc.",
    )
    declared_strategy: dict[str, Any] = Field(
        default_factory=dict,
        description="Active batting strategy",
    )
    bowling_strategy: dict[str, Any] = Field(
        default_factory=dict,
        description="Active bowling strategy",
    )
    field_setting: str = Field(default="Balanced")
    strategic_phase: str = Field(default="pre_ball")
    current_batter: dict[str, Any] = Field(default_factory=dict)
    non_striker: dict[str, Any] = Field(default_factory=dict)
    current_bowler: dict[str, Any] = Field(default_factory=dict)
    opponent_context: dict[str, Any] = Field(default_factory=dict)
    opponent_plan: dict[str, Any] = Field(default_factory=dict)
    last_outcome: dict[str, Any] = Field(default_factory=dict)
    eval_pack_id: str = Field(default="default")
    available_tools: list[str] = Field(default_factory=list)
    tool_history: list[dict[str, Any]] = Field(default_factory=list)
    last_ball_result: str = Field(default="")
    prompt_text: str = Field(default="")
    target: Optional[int] = Field(default=None)
    innings_type: str = Field(default="first")
    curriculum_stage: int = Field(default=2)
    match_plan: dict[str, Any] = Field(default_factory=dict)
    plan_review_due: bool = Field(default=False)
    plan_version: int = Field(default=0)
    plan_age_overs: int = Field(default=0)


class CricketState(State):
    """Episode state exposed via the /state endpoint."""

    game_state: str = Field(default="batting")
    phase: str = Field(default="powerplay")
    total_score: int = Field(default=0)
    wickets_lost: int = Field(default=0)
    over: int = Field(default=0)
    ball: int = Field(default=0)
    innings_type: str = Field(default="first")
    
    # Toss info
    toss_winner: Optional[str] = None
    toss_decision: Optional[str] = None
    
    # First innings stats (to set target for second)
    first_innings_score: int = 0
    target: Optional[int] = None
    match_result: str = Field(default="")
    reward_breakdown: dict[str, Any] = Field(default_factory=dict)
    innings_rewards: list[dict[str, Any]] = Field(default_factory=list)
    
    # Strategy and tool tracking
    coherence_scores: list[float] = Field(default_factory=list)
    adaptation_scores: list[float] = Field(default_factory=list)
    opponent_awareness_scores: list[float] = Field(default_factory=list)
    regret_scores: list[float] = Field(default_factory=list)
    tool_calls_made: int = Field(default=0)
    analyze_calls: list[dict[str, Any]] = Field(default_factory=list)
    tool_history: list[dict[str, Any]] = Field(default_factory=list)
    transcript: list[dict[str, Any]] = Field(default_factory=list)
    dls_par: float = Field(default=250.0)
    strategy_changes: int = Field(default=0)
    last_strategy_set_over: int = Field(default=-1)
    last_reflection: str = Field(default="")
    eval_pack_id: str = Field(default="default")
    opponent_mode: str = Field(default="heuristic")
    strategic_phase: str = Field(default="pre_ball")
    
    # Bowling specific
    bowling_strategy: dict[str, Any] = Field(default_factory=dict)
    field_setting: str = Field(default="Balanced")
    shot_plan: dict[str, Any] = Field(default_factory=dict)
    delivery_plan: dict[str, Any] = Field(default_factory=dict)
    opponent_plan: dict[str, Any] = Field(default_factory=dict)
    last_outcome: dict[str, Any] = Field(default_factory=dict)
    current_batter: dict[str, Any] = Field(default_factory=dict)
    non_striker: dict[str, Any] = Field(default_factory=dict)
    current_bowler: dict[str, Any] = Field(default_factory=dict)

    # Per-innings Dream11 batting stats (reset on innings change)
    batting_runs: int = Field(default=0)
    batting_balls_faced: int = Field(default=0)
    batting_fours: int = Field(default=0)
    batting_sixes: int = Field(default=0)
    batting_duck: bool = Field(default=False)

    # Per-innings Dream11 bowling stats (reset on innings change)
    bowling_wickets: int = Field(default=0)
    bowling_runs_conceded: int = Field(default=0)
    bowling_dot_balls: int = Field(default=0)
    bowling_lbw_bowled_wickets: int = Field(default=0)
    bowling_maiden_overs: int = Field(default=0)
    bowling_balls_in_over: int = Field(default=0)
    bowling_runs_in_over: int = Field(default=0)

    # Accumulated per-innings Dream11 scores (set at innings/match end)
    dream11_scores: list[float] = Field(default_factory=list)

    # Match plan state
    match_plan: dict[str, Any] = Field(default_factory=dict)
    plan_version: int = Field(default=0)
    plan_created_over: int = Field(default=-1)
    last_plan_update_over: int = Field(default=-1)
    plan_review_due: bool = Field(default=False)
    plan_staleness_penalties: list[float] = Field(default_factory=list)
    plan_commitment_scores: list[float] = Field(default_factory=list)
    plan_freshness_scores: list[float] = Field(default_factory=list)

    is_done: bool = Field(default=False)
    curriculum_stage: int = Field(default=2)
    max_overs: int = Field(default=50)