from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

import simpy

from kernbench.policy.routing.router import AddressResolver, PathRouter


@dataclass
class ComponentContext:
    """Topology services injected into every component implementation.

    Required by components that need routing or address resolution
    (IoCpuComponent, MCpuComponent, …).  TransitComponent ignores ctx.

    Passed via ComponentRegistry.create(node, overrides, ctx=ctx).
    """

    router: PathRouter
    resolver: AddressResolver
    positions: dict[str, tuple[float, float] | None]  # node_id → pos_mm
    ns_per_mm: float  # wire propagation constant (from topology spec)
    edge_map: dict[tuple[str, str], Any] = field(default_factory=dict)
    spec: dict = field(default_factory=dict)  # topology spec (cube layout, PE count, etc.)
    memory_store: Any = None  # MemoryStore for Phase 1 data-aware execution (ADR-0020)
    op_logger: Any = None     # OpLogger for Phase 1 op recording (ADR-0020)
    # node_id -> overhead_ns (ADR-0009 D5: used by M_CPU to compute per-PE
    # dispatch latency when stamping target_start_ns on KernelLaunchMsg).
    node_overhead_ns: dict[str, float] = field(default_factory=dict)

    def get_shared_resource(
        self, env: simpy.Environment, key: str, capacity: int = 1,
    ) -> simpy.Resource:
        """Return a shared SimPy Resource, creating it on first access.

        Used by PE components that share a resource across engines within
        the same PE (e.g. accel_slot shared by PE_GEMM and PE_MATH).
        Key should be scoped per PE: e.g. "sip0.cube0.pe0.accel_slot".
        """
        if not hasattr(self, "_shared_resources"):
            self._shared_resources: dict[str, simpy.Resource] = {}
        if key not in self._shared_resources:
            self._shared_resources[key] = simpy.Resource(env, capacity=capacity)
        return self._shared_resources[key]

    def compute_drain_ns(self, path: list[str], nbytes: int) -> float:
        """Wormhole drain time: nbytes / bottleneck_bw along path."""
        min_bw = float("inf")
        for i in range(len(path) - 1):
            edge = self.edge_map.get((path[i], path[i + 1]))
            if edge and getattr(edge, "bw_gbs", None):
                min_bw = min(min_bw, edge.bw_gbs)
        if min_bw == float("inf"):
            return 0.0
        return nbytes / min_bw

    def compute_path_latency_ns(self, path: list[str], nbytes: int = 0) -> float:
        """Formula latency along path: wire + per-node overhead + drain.

        ADR-0009 D5: M_CPU uses this to compute per-PE dispatch latency
        when stamping target_start_ns on KernelLaunchMsg fanout.
        """
        total = 0.0
        for i in range(len(path) - 1):
            edge = self.edge_map.get((path[i], path[i + 1]))
            if edge:
                total += edge.distance_mm * self.ns_per_mm
        for node_id in path:
            total += self.node_overhead_ns.get(node_id, 0.0)
        total += self.compute_drain_ns(path, nbytes)
        return total
