from pydantic import BaseModel
from typing import List, Tuple, Dict, Any, Optional
from enum import Enum


class Priority(str, Enum):
    CRITICAL = "critical"
    NORMAL = "normal"


class ActionType(str, Enum):
    MOVE = "move"
    ALLOCATE = "allocate"


class ServiceRequest(BaseModel):
    id: int
    position: Tuple[int, int]
    priority: Priority
    allocated: bool
    created_at: int


class Observation(BaseModel):
    agent_position: Tuple[int, int]
    grid: List[List[List[float]]]  # 5x5x3 grid
    requests: List[ServiceRequest]
    time_left: int
    resources_left: int
    step: int
    task_config: Dict[str, Any]


class Action(BaseModel):
    action_type: ActionType
    direction: Optional[str] = None  # For move actions


class Reward(BaseModel):
    score: float  # 0.0 to 1.0


class State(BaseModel):
    agent_position: Tuple[int, int]
    requests: List[Dict[str, Any]]
    time_left: int
    resources_left: int
    current_step: int
    done: bool
    total_reward: float