"""DMA Writeback Block: writes result tiles from TCM to HBM via real PE_DMA fabric.

Flow per tile:
  1. Receive flush Trigger from GEMM or MATH block
  2. Look up DmaWBDescriptor for address and tile size
  3. Issue DmaWriteCmd → PE_DMA → fabric → HBM controller → response
  4. Send completion Trigger to pipeline

Two _flush_loop processes run concurrently:
  - One drains GEMM → DMA_WB triggers (direct writeback path)
  - One drains MATH → DMA_WB triggers (K-accumulation or element-wise flush)
"""
from __future__ import annotations

import simpy

from kernbench.components.custom.pe_accel.types import DmaWBDescriptor, Trigger


class DmaWbBlock:
    """TCM → HBM tile writer. Shared across all concurrent pipelines.

    Pipelines pre-load DmaWBDescriptors keyed by (pipeline_id, tile_id).
    Each _flush_loop process reads triggers, issues real DmaWriteCmd to PE_DMA,
    and forwards completion to the pipeline's reply queue.
    """

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

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

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

    def _flush_loop(self, trig_q: simpy.Store):
        """Main process: receive flush triggers, issue DmaWriteCmd, send completion."""
        from kernbench.common.pe_commands import DmaWriteCmd, PeInternalTxn, TensorHandle

        while True:
            trigger: Trigger = yield trig_q.get()
            if trigger is None:
                break

            pid = trigger.pipeline_id
            tile_id = trigger.tile_id
            desc = self._descriptor_table.get((pid, tile_id))

            if desc:
                c_bytes = desc.Tm * desc.Tn * self._bpe

                # Issue real DMA write through PE_DMA → fabric → HBM
                write_done = self.env.event()
                handle = TensorHandle(
                    id=f"accel_wb_{pid}_{tile_id}",
                    addr=desc.dst_addr,
                    shape=(desc.Tm, desc.Tn),
                    dtype="float16",
                    nbytes=c_bytes,
                )
                txn = PeInternalTxn(
                    command=DmaWriteCmd(handle=handle, dst_addr=desc.dst_addr, nbytes=c_bytes),
                    done=write_done,
                    pe_prefix=self._pe_prefix,
                )
                t0 = self.env.now
                yield self._pe_dma_port.put(txn)
                yield write_done
                self.t_dma_write_per_tile.setdefault(pid, []).append(self.env.now - t0)

            yield self.completion_q.put(
                Trigger(tile_id=tile_id, pipeline_id=pid, source_block="DMA_WB")
            )
