"""DMA IN Block: reads tiles from HBM into TCM via real PE_DMA fabric.

Flow per tile:
  1. Receive DmaRequest from tiling pipeline
  2. Look up DmaInDescriptor for address and size
  3. Issue DmaReadCmd → PE_DMA → fabric → HBM controller → response
  4. Route completion Trigger to next block (GEMM, MATH, or COMPLETION)

Timing is real fabric latency (not analytical) — includes BW contention,
propagation delay, and HBM controller serialization.
"""
from __future__ import annotations

import simpy

from kernbench.components.custom.pe_accel.types import DmaInDescriptor, DmaRequest, Trigger


class DmaInBlock:
    """HBM → TCM tile reader. Shared across all concurrent pipelines.

    Pipelines pre-load DmaInDescriptors keyed by (pipeline_id, tile_id, operand).
    The _load_loop process reads DmaRequests, issues real DmaReadCmd to PE_DMA,
    and routes completion triggers based on descriptor.next_block.
    """

    def __init__(
        self,
        env: simpy.Environment,
        cmd_q: simpy.Store,
        to_fetch_trig: simpy.Store,
        to_math_trig: simpy.Store,
        completion_q: simpy.Store,
        *,
        pe_dma_port: simpy.Store | None,
        pe_prefix: str,
    ) -> None:
        self.env = env
        self.cmd_q = cmd_q
        self.to_fetch_trig = to_fetch_trig
        self.to_math_trig = to_math_trig
        self.completion_q = completion_q
        self._pe_dma_port = pe_dma_port
        self._pe_prefix = pe_prefix
        self._descriptor_table: dict[tuple[int, int, str], DmaInDescriptor] = {}

        # Per-pipeline timing histogram (keyed by pipeline_id)
        self.t_dma_read_per_request: dict[int, list[float]] = {}

    def load_descriptors(self, descs: dict[tuple, DmaInDescriptor]) -> None:
        """Pre-load per-operand DMA descriptors (cumulative across pipelines)."""
        self._descriptor_table.update(descs)

    def _load_loop(self):
        """Main process: receive DmaRequests, issue DmaReadCmd, route triggers."""
        from kernbench.common.pe_commands import DmaReadCmd, PeInternalTxn, TensorHandle

        while True:
            req: DmaRequest = yield self.cmd_q.get()
            if req is None:
                break

            desc = self._descriptor_table[(req.pipeline_id, req.tile_id, req.operand)]

            # Issue real DMA read through PE_DMA → fabric → HBM
            read_done = self.env.event()
            handle = TensorHandle(
                id=f"accel_rd_{req.pipeline_id}_{req.tile_id}_{req.operand}",
                addr=desc.src_addr,
                shape=(desc.size_bytes,),
                dtype="uint8",
                nbytes=desc.size_bytes,
            )
            txn = PeInternalTxn(
                command=DmaReadCmd(handle=handle, src_addr=desc.src_addr, nbytes=desc.size_bytes),
                done=read_done,
                pe_prefix=self._pe_prefix,
            )
            t0 = self.env.now
            yield self._pe_dma_port.put(txn)
            yield read_done
            self.t_dma_read_per_request.setdefault(req.pipeline_id, []).append(self.env.now - t0)

            # Route trigger to next block
            trig = Trigger(
                tile_id=req.tile_id,
                pipeline_id=req.pipeline_id,
                vc=0 if req.operand == "A" else 1,
                source_block="DMA_IN",
            )
            if desc.next_block == "MATH":
                yield self.to_math_trig.put(trig)
            elif desc.next_block == "COMPLETION":
                yield self.completion_q.put(trig)
            else:  # "GEMM" (default)
                yield self.to_fetch_trig.put(trig)
