backend/app/models/requests.py

"""Pydantic models for API requests and responses."""

from pydantic import BaseModel, Field
from typing import Optional, List, Tuple
from enum import Enum


class Algorithm(str, Enum):
    """Available search algorithms."""

    BF = "BF"  # Breadth-first search
    DF = "DF"  # Depth-first search
    ID = "ID"  # Iterative deepening
    UC = "UC"  # Uniform cost search
    GR1 = "GR1"  # Greedy with Manhattan heuristic
    GR2 = "GR2"  # Greedy with Euclidean heuristic
    AS1 = "AS1"  # A* with Manhattan heuristic
    AS2 = "AS2"  # A* with Tunnel-aware heuristic


class Position(BaseModel):
    """A position on the grid."""

    x: int
    y: int

    def to_tuple(self) -> Tuple[int, int]:
        return (self.x, self.y)


class SegmentData(BaseModel):
    """Segment data for API."""

    src: Position
    dst: Position
    traffic: int = Field(ge=0, le=4)


class StoreData(BaseModel):
    """Store data for API."""

    id: int
    position: Position


class DestinationData(BaseModel):
    """Destination data for API."""

    id: int
    position: Position


class TunnelData(BaseModel):
    """Tunnel data for API."""

    entrance1: Position
    entrance2: Position
    cost: Optional[int] = None


# Request Models


class GridConfig(BaseModel):
    """Configuration for grid generation."""

    width: Optional[int] = Field(None, ge=5, le=50)
    height: Optional[int] = Field(None, ge=5, le=50)
    num_stores: Optional[int] = Field(None, ge=1, le=3)
    num_destinations: Optional[int] = Field(None, ge=1, le=10)
    num_tunnels: Optional[int] = Field(None, ge=0, le=10)
    obstacle_density: float = Field(0.1, ge=0.0, le=0.5)


class SearchRequest(BaseModel):
    """Request for running a search/plan."""

    initial_state: str
    traffic: str
    strategy: Algorithm
    visualize: bool = False


class PathRequest(BaseModel):
    """Request for finding a single path."""

    grid_width: int
    grid_height: int
    start: Position
    goal: Position
    segments: List[SegmentData]
    tunnels: List[TunnelData] = []
    strategy: Algorithm


class CompareRequest(BaseModel):
    """Request for comparing all algorithms."""

    initial_state: str
    traffic: str


# Response Models


class PathData(BaseModel):
    """Path result data."""

    plan: str
    cost: float
    nodes_expanded: int
    path: List[Position]


class GridData(BaseModel):
    """Complete grid state data."""

    width: int
    height: int
    stores: List[StoreData]
    destinations: List[DestinationData]
    tunnels: List[TunnelData]
    segments: List[SegmentData]


class GenerateResponse(BaseModel):
    """Response from grid generation."""

    initial_state: str
    traffic: str
    parsed: GridData


class SearchResponse(BaseModel):
    """Response from search/plan execution."""

    plan: str
    cost: float
    nodes_expanded: int
    runtime_ms: float
    memory_kb: float
    cpu_percent: float
    path: List[Position]
    steps: Optional[List[dict]] = None


class PlanResponse(BaseModel):
    """Response from delivery planning."""

    output: str
    assignments: List[dict]
    total_cost: float
    total_nodes_expanded: int
    runtime_ms: float
    memory_kb: float
    cpu_percent: float


class ComparisonResult(BaseModel):
    """Result of comparing a single algorithm."""

    algorithm: str
    name: str
    plan: str
    cost: float
    nodes_expanded: int
    runtime_ms: float
    memory_kb: float
    cpu_percent: float
    is_optimal: bool = False


class CompareResponse(BaseModel):
    """Response from algorithm comparison."""

    comparisons: List[ComparisonResult]
    optimal_cost: float


class AlgorithmInfo(BaseModel):
    """Information about an algorithm."""

    code: str
    name: str
    description: str


class AlgorithmsResponse(BaseModel):
    """List of available algorithms."""

    algorithms: List[AlgorithmInfo]