"""Grid and Segment models for the delivery search problem."""

from dataclasses import dataclass, field
from typing import Dict, Tuple, Optional


@dataclass
class Segment:
    """Represents a road segment between two adjacent grid points."""

    src: Tuple[int, int]
    dst: Tuple[int, int]
    traffic: int  # 0 = blocked, 1-4 = traffic level

    def __post_init__(self):
        # Normalize segment direction (ensure src < dst lexicographically)
        if self.src > self.dst:
            self.src, self.dst = self.dst, self.src

    @property
    def is_blocked(self) -> bool:
        return self.traffic == 0

    def get_key(self) -> Tuple[Tuple[int, int], Tuple[int, int]]:
        """Get normalized key for segment lookup."""
        return (self.src, self.dst)


@dataclass
class Grid:
    """Represents the city grid with all road segments."""

    width: int
    height: int
    segments: Dict[Tuple[Tuple[int, int], Tuple[int, int]], Segment] = field(
        default_factory=dict
    )

    def get_segment(
        self, src: Tuple[int, int], dst: Tuple[int, int]
    ) -> Optional[Segment]:
        """Get segment between two points (order doesn't matter)."""
        key = (src, dst) if src < dst else (dst, src)
        return self.segments.get(key)

    def get_traffic(self, src: Tuple[int, int], dst: Tuple[int, int]) -> int:
        """Get traffic level for segment between two points."""
        segment = self.get_segment(src, dst)
        return segment.traffic if segment else 0

    def is_blocked(self, src: Tuple[int, int], dst: Tuple[int, int]) -> bool:
        """Check if segment between two points is blocked."""
        return self.get_traffic(src, dst) == 0

    def is_valid_position(self, pos: Tuple[int, int]) -> bool:
        """Check if position is within grid bounds."""
        x, y = pos
        return 0 <= x < self.width and 0 <= y < self.height

    def add_segment(self, src: Tuple[int, int], dst: Tuple[int, int], traffic: int):
        """Add or update a segment."""
        segment = Segment(src, dst, traffic)
        self.segments[segment.get_key()] = segment

    def get_neighbors(self, pos: Tuple[int, int]) -> list[Tuple[int, int]]:
        """Get all valid neighboring positions (not blocked)."""
        x, y = pos
        neighbors = []
        directions = [(0, 1), (0, -1), (1, 0), (-1, 0)]  # up, down, right, left

        for dx, dy in directions:
            new_pos = (x + dx, y + dy)
            if self.is_valid_position(new_pos) and not self.is_blocked(pos, new_pos):
                neighbors.append(new_pos)

        return neighbors

    def to_dict(self) -> dict:
        """Convert grid to dictionary for JSON serialization."""
        return {
            "width": self.width,
            "height": self.height,
            "segments": [
                {
                    "src": {"x": seg.src[0], "y": seg.src[1]},
                    "dst": {"x": seg.dst[0], "y": seg.dst[1]},
                    "traffic": seg.traffic,
                }
                for seg in self.segments.values()
            ],
        }