from __future__ import annotations

from collections.abc import Generator
from typing import TYPE_CHECKING, Any

import simpy

from kernbench.components.base import ComponentBase
from kernbench.sim_engine.transaction import Flit

if TYPE_CHECKING:
    from kernbench.components.context import ComponentContext
    from kernbench.topology.types import Node


class TransitComponent(ComponentBase):
    """Transit component for NOC, UCIe, XBAR nodes (ADR-0033 Phase 2c).

    Flit-aware pass-through: forwards each Flit to the next hop with
    per-transaction ``overhead_ns`` applied ONCE (at first-flit arrival,
    modeling header decode + routing decision). Subsequent flits of the
    same transaction pipeline through with no extra delay, preserving
    wormhole-style cut-through across multi-hop paths.

    Forwarding is SERIAL in the worker: each flit is forwarded in arrival
    order. Spawning ``env.process`` per flit would let later flits
    overtake earlier ones (when the first flit yields ``overhead_ns``
    while subsequent flits skip it), producing out-of-order delivery
    and early ``is_last`` signaling at the destination.

    Non-Flit messages (zero-byte control Transactions, etc.) fall back
    to the legacy atomic ``_forward_txn`` path via ``env.process``.
    """

    _FLIT_AWARE = True

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)
        self._txn_decoded: set[int] = set()

    def run(self, env: simpy.Environment, nbytes: int) -> Generator:
        overhead_ns = float(self.node.attrs.get("overhead_ns", 0.0))
        yield env.timeout(overhead_ns)

    def _worker(self, env: simpy.Environment) -> Generator:
        while True:
            msg: Any = yield self._inbox.get()
            if isinstance(msg, Flit):
                tid = id(msg.txn)
                if tid not in self._txn_decoded:
                    self._txn_decoded.add(tid)
                    yield from self.run(env, msg.txn.nbytes)
                if msg.is_last:
                    self._txn_decoded.discard(tid)
                next_hop = self._next_hop_in_path(msg.txn)
                if next_hop and next_hop in self.out_ports:
                    yield self.out_ports[next_hop].put(msg)
                elif msg.is_last:
                    msg.txn.done.succeed()
            else:
                env.process(self._forward_txn(env, msg))

    def _next_hop_in_path(self, txn: Any) -> str | None:
        my_id = self.node.id
        path = getattr(txn, "path", None)
        if not path:
            return None
        for i, n in enumerate(path):
            if n == my_id and i + 1 < len(path):
                return path[i + 1]
        return None
