from __future__ import annotations

from collections.abc import Generator
from typing import TYPE_CHECKING

from kernbench.components.base import ComponentBase

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


class PeTcmComponent(ComponentBase):
    """PE_TCM: tightly-coupled memory / local SRAM staging buffer.

    Terminal storage component for PE-internal dataflow (ADR-0014 D5).
    Phase 0: applies overhead_ns and drain_ns at terminal.
    """

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)

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