Network Object-Oriented Design

from dataclasses import dataclass, field
from typing import Optional
import heapq

@dataclass
class Link:
    node_a: "Node"
    node_b: "Node"
    bandwidth_mbps: float
    latency_ms: float
    is_up: bool = True

    def other(self, node: "Node") -> "Node":
        return self.node_b if node is self.node_a else self.node_a

@dataclass
class Node:
    node_id: str
    ip_address: str
    links: list[Link] = field(default_factory=list)

    def connect(self, link: Link) -> None:
        self.links.append(link)

    def active_links(self) -> list[Link]:
        return [l for l in self.links if l.is_up]

@dataclass
class Packet:
    source_ip: str
    destination_ip: str
    payload: bytes
    ttl: int = 64


class Network:
    def __init__(self):
        self._nodes: dict[str, Node] = {}

    def add_node(self, node_id: str, ip_address: str) -> Node:
        node = Node(node_id=node_id, ip_address=ip_address)
        self._nodes[node_id] = node
        return node

    def add_link(self, id_a: str, id_b: str,
                 bandwidth_mbps: float, latency_ms: float) -> Link:
        node_a = self._nodes[id_a]
        node_b = self._nodes[id_b]
        link = Link(node_a, node_b, bandwidth_mbps, latency_ms)
        node_a.connect(link)
        node_b.connect(link)
        return link

    def shortest_path(self, src_id: str, dst_id: str) -> Optional[list[str]]:
        """Dijkstra by latency. Returns list of node IDs or None if unreachable."""
        dist = {nid: float("inf") for nid in self._nodes}
        prev: dict[str, Optional[str]] = {nid: None for nid in self._nodes}
        dist[src_id] = 0
        heap = [(0.0, src_id)]

        while heap:
            d, uid = heapq.heappop(heap)
            if d > dist[uid]:
                continue
            if uid == dst_id:
                break
            for link in self._nodes[uid].active_links():
                neighbor = link.other(self._nodes[uid]).node_id
                new_dist = d + link.latency_ms
                if new_dist < dist[neighbor]:
                    dist[neighbor] = new_dist
                    prev[neighbor] = uid
                    heapq.heappush(heap, (new_dist, neighbor))

        if dist[dst_id] == float("inf"):
            return None

        path, cur = [], dst_id
        while cur:
            path.append(cur)
            cur = prev[cur]
        return list(reversed(path))

    def send_packet(self, packet: Packet, src_id: str, dst_id: str) -> bool:
        """Returns True if packet reaches destination."""
        path = self.shortest_path(src_id, dst_id)
        if not path:
            return False
        print(f"Packet routed: {' → '.join(path)}")
        return True