"""PE_FETCH_STORE: TCM ↔ Register File transfer unit (ADR-0014 D1).

Handles both fetch (TCM → register) and store (register → TCM).
BW serialization is delegated to PE_TCM via port communication.
"""
from __future__ import annotations

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

import simpy

from kernbench.components.base import PeEngineBase

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


class PeFetchStoreComponent(PeEngineBase):
    """PE_FETCH_STORE: TCM ↔ Register File (ADR-0014 D1).

    Receives TileTokens via pipeline self-routing.
    Sends TcmRequest to PE_TCM for BW-based latency.
    """

    def __init__(self, node: Node, ctx: ComponentContext | None = None) -> None:
        super().__init__(node, ctx)
        self._tcm_id = f"{self._pe_prefix}.pe_tcm"

    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:
        """Handle both PeInternalTxn (legacy) and TileToken (pipeline)."""
        from kernbench.common.pe_commands import PeInternalTxn
        from kernbench.components.builtin.pe_types import TileToken

        while True:
            msg: Any = yield self._inbox.get()
            if isinstance(msg, TileToken):
                env.process(self._pipeline_process(env, msg))
            elif isinstance(msg, PeInternalTxn):
                env.process(self.handle_command(env, msg))
            else:
                env.process(self._forward_txn(env, msg))

    def _pipeline_process(self, env: simpy.Environment, token: Any) -> Generator:
        """Process a pipeline TileToken: fetch or store via TCM."""
        from kernbench.components.builtin.pe_tcm import TcmRequest

        self._on_process_start(env, token)

        direction = token.params.get("direction", "read")
        nbytes = token.params.get("nbytes", 0)

        if nbytes > 0 and self._tcm_id in self.out_ports:
            done = env.event()
            yield self.out_ports[self._tcm_id].put(
                TcmRequest(direction=direction, nbytes=nbytes, done=done)
            )
            yield done

        self._on_process_end(env, token)

        # Self-routing: advance to next stage
        next_stage = token.advance()
        if next_stage is not None:
            yield self.out_ports[next_stage.component].put(token)
        else:
            token.pipeline_ctx.complete_tile()

    def handle_command(self, env: simpy.Environment, pe_txn: Any) -> Generator:
        """Legacy PeInternalTxn handling."""
        yield from self.run(env, 0)
        pe_txn.done.succeed()
