"""
power.py — Capacity-aware power flow simulation for Powernet.

Algorithm:
  1. BFS from all generators to identify which buildings are reachable.
  2. For each connected component that contains ≥1 generator:
     a. Sum total generation MW.
     b. Sort consumers/research by BFS distance from nearest generator.
     c. Allocate generation to consumers in order (closest first).
     d. Consumers that exceed remaining capacity are unpowered.
  3. Wire-tier capacity: a single wire limits how much MW can flow through it.
     If the aggregated power demand routed through a wire would exceed its
     tier capacity, that wire is "overloaded" and the consumer on its far
     side is unpowered. (Simplified: we check per-wire capacity as part of
     the BFS path selection.)
  4. Generator overload: if total island demand > total island generation,
     mark all generators in that island as overloaded.
  5. Wire flows: for each wire, estimate MW flowing through it based on
     powered consumers reachable only via that wire from the generator side.
"""

from __future__ import annotations

from collections import deque
from typing import TYPE_CHECKING

from helpers.game_data import BUILDING_TYPES
from .constants import WIRE_TIERS

if TYPE_CHECKING:
    from helpers.game_data import GameState


def compute_powered(
    gs: "GameState",
) -> tuple[set[str], dict[str, float], dict[str, float]]:
    """
    Capacity-aware BFS power simulation.

    Returns:
        powered       — set of build_ids that receive power
        overloaded    — dict mapping generator build_id → overload_mw
                        (generators in islands where demand > supply)
        wire_flows    — dict mapping wire key (from_id + "_" + to_id) → MW flowing
    """
    # ---- Build adjacency: bid → list[(neighbour_bid, wire_layer, wire_cap_mw)] ----
    adj: dict[str, list[tuple[str, int, float]]] = {}
    for b in gs.buildings:
        adj[b.build_id] = []

    for w in gs.wires:
        if w.from_id not in adj or w.to_id not in adj:
            continue
        tier    = w.layer
        cap_mw  = WIRE_TIERS[tier][1] if tier < len(WIRE_TIERS) else WIRE_TIERS[-1][1]
        adj[w.from_id].append((w.to_id,   tier, cap_mw))
        adj[w.to_id  ].append((w.from_id, tier, cap_mw))

    # ---- Find all connected components ----
    visited: set[str] = set()
    components: list[list[str]] = []

    for start in list(adj.keys()):
        if start in visited:
            continue
        component: list[str] = []
        q: deque[str] = deque([start])
        visited.add(start)
        while q:
            cur = q.popleft()
            component.append(cur)
            for (nb, *_) in adj.get(cur, []):
                if nb not in visited:
                    visited.add(nb)
                    q.append(nb)
        components.append(component)

    powered:     set[str]          = set()
    overloaded:  dict[str, float]  = {}   # generator build_id → overload_mw
    wire_flows:  dict[str, float]  = {}   # "from_id_to_id" → MW

    for component in components:
        bids = set(component)

        # Sum generation in this component
        total_gen: float = 0.0
        generators: list[str] = []
        for bid in bids:
            b  = gs.get_building_by_id(bid)
            td = BUILDING_TYPES.get(b.type_id) if b else None
            if td and td.power_mw > 0:
                total_gen += td.power_mw
                generators.append(bid)

        if not generators:
            # No generators → nothing is powered in this island
            continue

        # Sum total demand in island to detect overload
        total_demand: float = 0.0
        for bid in bids:
            if bid in generators:
                continue
            b  = gs.get_building_by_id(bid)
            td = BUILDING_TYPES.get(b.type_id) if b else None
            if td and td.power_mw < 0:
                total_demand += abs(td.power_mw)

        # Determine island overload
        island_overloaded = total_demand > total_gen
        overload_amount   = max(0.0, total_demand - total_gen)

        if island_overloaded:
            # Mark all generators in island as overloaded; they output 0
            for gid in generators:
                overloaded[gid] = overload_amount
            # Overloaded island: generators show as powered but consumers are not
            for gid in generators:
                powered.add(gid)
            continue

        # Island has enough generation — proceed with capacity-aware allocation
        for gid in generators:
            powered.add(gid)

        # BFS from generators; track distance from nearest generator
        dist: dict[str, int] = {}
        via_wire: dict[str, tuple[str, str, int]] = {}   # bid → (from_id, to_id, layer)
        bfs_q: deque[str] = deque()
        for gid in generators:
            dist[gid] = 0
            bfs_q.append(gid)

        # Track min wire cap on path from generator to each node
        path_cap: dict[str, float] = {gid: 1e18 for gid in generators}

        while bfs_q:
            cur = bfs_q.popleft()
            for (nb, tier, cap_mw) in adj.get(cur, []):
                if nb not in dist:
                    dist[nb] = dist[cur] + 1
                    # Track bottleneck capacity on path
                    path_cap[nb] = min(path_cap.get(cur, 1e18), cap_mw)
                    # Record which wire was used to reach nb
                    via_wire[nb] = (cur, nb, tier)
                    bfs_q.append(nb)

        # Sort consumers/research by BFS distance (closest first)
        consumers: list[tuple[int, str, float]] = []
        for bid in bids:
            if bid in generators:
                continue
            b  = gs.get_building_by_id(bid)
            td = BUILDING_TYPES.get(b.type_id) if b else None
            if td is None or bid not in dist:
                continue
            demand = abs(td.power_mw) if td.power_mw < 0 else 0.0
            # Research buildings also consume power
            if td.power_mw < 0:
                consumers.append((dist[bid], bid, demand))

        consumers.sort(key=lambda x: x[0])

        # Allocate generation capacity to consumers in order
        remaining = total_gen
        wire_load: dict[tuple[str, str, int], float] = {}   # (f,t,layer) → MW routed

        for _, bid, demand in consumers:
            # Check bottleneck wire capacity on path
            cap_on_path = path_cap.get(bid, 0.0)
            if cap_on_path <= 0:
                continue

            if remaining <= 0:
                break

            if demand <= remaining and demand <= cap_on_path:
                powered.add(bid)
                remaining -= demand
                # Accumulate load on the wire that connects this building
                wire_key = via_wire.get(bid)
                if wire_key:
                    wire_load[wire_key] = wire_load.get(wire_key, 0.0) + demand

        # Build wire_flows from wire_load
        # For each wire that carried load, propagate flow upstream too.
        # Simple approach: for each consumer's via_wire, record flow.
        # We also record the demand that *passes through* each wire by
        # accumulating downstream demand: BFS-reachable powered consumers
        # from each wire's far side.
        #
        # Efficient approach: since we have via_wire (shortest-path tree),
        # we can accumulate flows up the tree.
        # child_flow[bid] = demand of powered consumers in subtree rooted at bid.
        child_flow: dict[str, float] = {}
        for _, bid, demand in consumers:
            if bid in powered:
                child_flow[bid] = demand

        # Walk from leaves to root accumulating flow on each edge
        # Process in reverse BFS order (furthest first)
        ordered_bids = sorted(
            [bid for bid in bids if bid not in generators and bid in dist],
            key=lambda b2: -dist[b2]
        )
        for bid in ordered_bids:
            flow_here = child_flow.get(bid, 0.0)
            if flow_here <= 0:
                continue
            wire_key = via_wire.get(bid)
            if wire_key:
                fid, tid_, _ = wire_key
                k1 = fid + "_" + tid_
                k2 = tid_ + "_" + fid
                wire_flows[k1] = wire_flows.get(k1, 0.0) + flow_here
                wire_flows[k2] = wire_flows.get(k2, 0.0) + flow_here
                # Propagate up to parent
                child_flow[fid] = child_flow.get(fid, 0.0) + flow_here

    return powered, overloaded, wire_flows
