module.py

from typing import TypedDict
from rich.console import Console


class DictMetaDataZone(TypedDict):
    """Typed metadata stored on a zone."""

    zone: str
    color: str
    max_drones: int | float


class DictMetaDataConnection(TypedDict):
    """Typed metadata stored on a connection."""

    max_link_capacity: int


class Zone:
    def __init__(
        self, name: str, coordinates: tuple[int, int], metadata: str | None
    ) -> None:
        """Create a zone with optional metadata.

        Args:
            name: Zone name.
            coordinates: Zone coordinates on the map.
            metadata: Optional metadata string in ``key=value`` form.
        """
        self.name = name
        self.coordinates = coordinates
        self.default_metadata: DictMetaDataZone = {
            "zone": "normal",
            "color": "none",
            "max_drones": 1,
        }

        if metadata is not None and metadata.strip():
            self.set_metadata(metadata)

        self.available = self.default_metadata["max_drones"]

        if self.default_metadata["zone"] == "priority":
            self.priority = 0
        else:
            self.priority = 1

    def set_metadata(self, metadata: str) -> None:
        """Parse and validate zone metadata.

        Args:
            metadata: Metadata string in ``key=value`` form.

        Raises:
            ValueError: If the metadata format or values are invalid.
        """
        zone_types = ["normal", "blocked", "restricted", "priority"]
        colors = [
            "black",
            "blue",
            "brown",
            "crimson",
            "cyan",
            "darkred",
            "gold",
            "green",
            "lime",
            "magenta",
            "maroon",
            "orange",
            "purple",
            "rainbow",
            "red",
            "violet",
            "yellow",
            "gray"
        ]

        tags = metadata.split()
        for tag in tags:
            try:
                key, value = tag.split("=")
            except ValueError:
                raise ValueError(
                    "Invalid zone metadata: expected key=value"
                )
            if key == "zone":
                if value not in zone_types:
                    raise ValueError(
                        "Invalid zone metadata for 'zone'"
                    )
                self.default_metadata["zone"] = value
            elif key == "color":
                if value not in colors:
                    raise ValueError(
                        "Invalid zone metadata for 'color'"
                    )
                self.default_metadata["color"] = value
            elif key == "max_drones":
                try:
                    max_drones = int(value)
                except ValueError:
                    raise ValueError(
                        "Invalid zone metadata for 'max_drones': "
                        "expected a positive integer"
                    )
                if max_drones <= 0:
                    raise ValueError(
                        "Invalid zone metadata for 'max_drones': "
                        "must be greater than 0"
                    )
                self.default_metadata["max_drones"] = max_drones
            else:
                raise ValueError("Invalid zone metadata key")


class StartHub(Zone):
    def __init__(
        self, name: str, coordinates: tuple[int, int], metadata: str | None
    ) -> None:
        """Create the starting hub.

        Args:
            name: Hub name.
            coordinates: Hub coordinates on the map.
            metadata: Optional metadata string in ``key=value`` form.
        """
        super().__init__(name, coordinates, metadata)
        self.available = float('inf')


class EndHub(Zone):
    def __init__(
        self, name: str, coordinates: tuple[int, int], metadata: str | None
    ) -> None:
        """Create the ending hub.

        Args:
            name: Hub name.
            coordinates: Hub coordinates on the map.
            metadata: Optional metadata string in ``key=value`` form.
        """
        super().__init__(name, coordinates, metadata)
        self.available = float('inf')


class Hub(Zone):
    def __init__(
        self, name: str, coordinates: tuple[int, int], metadata: str | None
    ) -> None:
        """Create an intermediate hub.

        Args:
            name: Hub name.
            coordinates: Hub coordinates on the map.
            metadata: Optional metadata string in ``key=value`` form.
        """
        super().__init__(name, coordinates, metadata)


class Connection:
    def __init__(self, zone1: Zone, zone2: Zone, metadata: str | None) -> None:
        """Create a connection between two zones.

        Args:
            zone1: First endpoint.
            zone2: Second endpoint.
            metadata: Optional metadata string in ``key=value`` form.
        """
        self.zone1 = zone1
        self.zone2 = zone2
        self.default_metadata: DictMetaDataConnection = {
            "max_link_capacity": 1
        }
        self.fallback_max_link_capacity = 1

        if metadata is not None and metadata.strip():
            self.set_metadata(metadata)

    def set_metadata(self, metadata: str) -> None:
        """Parse and validate connection metadata.

        Args:
            metadata: Metadata string in ``key=value`` form.

        Raises:
            ValueError: If the metadata format or values are invalid.
        """
        tags = metadata.split()
        for tag in tags:
            try:
                key, value = tag.split("=")
            except ValueError:
                raise ValueError(
                    "Invalid connection metadata: expected key=value"
                )
            if key != "max_link_capacity":
                raise ValueError("Invalid connection metadata key")
            try:
                max_link_capacity = int(value)
            except ValueError:
                raise ValueError(
                    "Invalid connection metadata for 'max_link_capacity': "
                    "expected as integer"
                )
            if max_link_capacity <= 0:
                raise ValueError(
                    "Invalid connection metadata for 'max_link_capacity': "
                    "must be greater than 0"
                )

        self.default_metadata["max_link_capacity"] = max_link_capacity
        self.fallback_max_link_capacity = max_link_capacity


class ZoneFactory:
    """Factory for building zone instances from parsed map data."""

    @staticmethod
    def create(
        create_type: str, name: str,
        coordinates: tuple[int, int], metadata: str
    ) -> Zone:
        """Create a zone instance of the requested type.

        Args:
            create_type: Zone kind, such as ``start``, ``end`` or ``hub``.
            name: Zone name.
            coordinates: Zone coordinates on the map.
            metadata: Optional metadata string.

        Returns:
            A zone instance matching ``create_type``.

        Raises:
            ValueError: If ``create_type`` is not supported.
        """
        match create_type:
            case "start":
                return StartHub(name, coordinates, metadata)
            case "end":
                return EndHub(name, coordinates, metadata)
            case "hub":
                return Hub(name, coordinates, metadata)
            case _:
                raise ValueError(
                    "Invalid zone type: expected start, end, or hub"
                )


class Drone:
    def __init__(self, id: str, position: Zone, path: list[str]) -> None:
        """Create a drone with its current position and planned path.

        Args:
            id: Drone identifier.
            position: Current zone.
            path: Remaining path as zone names.
        """
        self.id = id
        self.position = position
        self.path = path
        self.transit_turns_left = 0
        self.next_zone = position


class Graph:
    """Weighted graph of zones and connections."""

    zone_cost = {
        "normal": 1,
        "blocked": float('inf'),
        "restricted": 2,
        "priority": 1
    }

    def __init__(self) -> None:
        """Initialize an empty adjacency list."""
        self.adj_list: dict[Zone, list[tuple[int | float, Zone]]] = dict()

    def add_nodes(self, node: Zone) -> None:
        """Add a zone to the graph.

        Args:
            node: Zone to add.

        Raises:
            ValueError: If the node already exists.
        """
        if node not in self.adj_list:
            self.adj_list[node] = list()
        else:
            raise ValueError("Node exist already")

    def add_edges(self, connection: Connection) -> None:
        """Add a bidirectional edge for a connection.

        Args:
            connection: Connection to register in the graph.

        Raises:
            ValueError: If the edge already exists.
        """
        zone1_cost = self.zone_cost[connection.zone1.default_metadata["zone"]]
        zone2_cost = self.zone_cost[connection.zone2.default_metadata["zone"]]

        if ((zone1_cost, connection.zone1) not in
            self.adj_list[connection.zone2]

           and (zone2_cost, connection.zone2) not in
           self.adj_list[connection.zone1]):

            self.adj_list[connection.zone2].append((
                zone1_cost,
                connection.zone1
            ))
            self.adj_list[connection.zone1].append((
                zone2_cost,
                connection.zone2
            ))
        else:
            raise ValueError("Edge exist already")

    def get_neighbors(self, node: Zone) -> list[tuple[int | float, Zone]]:
        """Return adjacent zones for a node.

        Args:
            node: Zone whose neighbors are requested.

        Returns:
            A list of ``(cost, zone)`` neighbor pairs.
        """
        return self.adj_list[node]


class Visualization:
    """Format zones for colored terminal output."""

    COLORS = {
        "none": "#FFFFFF",
        "black": "#000000",
        "blue": "#0000FF",
        "brown": "#A52A2A",
        "crimson": "#DC143C",
        "cyan": "#00FFFF",
        "darkred": "#8B0000",
        "gold": "#FFD700",
        "green": "#008000",
        "lime": "#00FF00",
        "magenta": "#FF00FF",
        "maroon": "#800000",
        "orange": "#FFA500",
        "purple": "#800080",
        "red": "#FF0000",
        "violet": "#EE82EE",
        "yellow": "#FFFF00",
        "gray": "#808080"
    }

    Console = Console

    @classmethod
    def apply_color(cls, zone: Zone) -> str:
        """Return a rich-formatted zone name.

        Args:
            zone: Zone to format.

        Returns:
            The zone name wrapped in color markup.
        """
        if zone.default_metadata["color"] == "rainbow":
            name = ""
            for i, char in enumerate(zone.name, 0):
                color = list(cls.COLORS.values())[i % len(cls.COLORS)]
                name += f"[{color}]{char}"

            return name

        return f"[{cls.COLORS[zone.default_metadata['color']]}]{zone.name}"