import heapq


class SimulationWorld:

    def __init__(self, engine, num_nodes):
        self.engine = engine
        self.num_nodes = num_nodes
        self.rescued_pool = 0

    def get_node_state(self, node_id):
        node = self.engine.get_node_state(node_id)
        if node_id == 0:
            node.survivors = self.rescued_pool
            node.temperature = 20.0
            node.water_level = 0.0
            node.is_collapsed = False
            node.stability = 1.0
            node.status.infra_health = 1.0
            node.pvi = 0.0
            return node

        return node

    def get_neighbors(self, node_id):
        return self.engine.get_neighbors(node_id)

    def allocate_resources(self, node_id, resource, amount):
        self.engine.allocate_resources(node_id, resource, amount)

    def clear_debris(self, node_id, amount):
        self.engine.clear_debris(node_id, amount)

    def apply_impact(self, node_id, t_delta, w_delta, p_delta):
        self.engine.apply_impact(node_id, t_delta, w_delta, p_delta)

    def step(self):
        self.engine.step()

    def notify_extraction(self, node_id, count):
        if node_id == 0:
            return
        actual = self.engine.decrement_survivors(node_id, count)
        if actual > 0:
            self.rescued_pool += actual

    def notify_evacuation(self, node_id, count):
        self.notify_extraction(node_id, count)

    def transfer_survivors(self, source_id, target_id, count):
        if source_id == 0:
            source_count = self.rescued_pool
        else:
            source_count = self.engine.get_node_state(source_id).survivors

        actual = min(count, source_count)

        if source_id == 0:
            self.rescued_pool -= actual
        else:
            self.engine.decrement_survivors(source_id, actual)

        if target_id == 0:
            self.rescued_pool += actual
        else:
            self.engine.increment_survivors(target_id, actual)

        return actual

    def get_snapshot(self):
        return {
            "engine_snapshot": self.engine.get_snapshot(),
            "rescued_pool": self.rescued_pool,
            "num_nodes": self.num_nodes
        }

    def restore_snapshot(self, snapshot):
        self.engine.restore_snapshot(snapshot["engine_snapshot"])
        self.rescued_pool = snapshot["rescued_pool"]
        self.num_nodes = snapshot["num_nodes"]

    def next_hop(self, start: int, target: int):
        if start == target:
            return start
        dist = {start: 0.0}
        parent = {start: None}
        heap = [(0.0, start)]
        while heap:
            cost, node = heapq.heappop(heap)
            if node == target:
                break
            if cost > dist.get(node, float("inf")):
                continue
            for nb in self.get_neighbors(node):
                nb_state = self.get_node_state(nb)
                if nb_state.is_collapsed or nb_state.water_level > 7:
                    continue

                edge_w = 1.0 + (1.0 - nb_state.status.infra_health) * 2.0
                if nb == 0 and target != 0:
                    edge_w += 10.0

                new_cost = cost + edge_w
                if new_cost < dist.get(nb, float("inf")):
                    dist[nb] = new_cost
                    parent[nb] = node
                    heapq.heappush(heap, (new_cost, nb))
        if target not in parent:
            return None
        node = target
        while parent[node] != start:
            node = parent[node]
        return node