"""TCM Block: tightly-coupled memory with BW-based access serialization.

Models SPMem (scratchpad memory) inside the PE. Compute blocks (GEMM, MATH)
send TcmRequests for load/store operations. The TCM block serializes access
per channel and computes timing based on data size and bandwidth.

Two channels:
  - READ  (SPMem → compute unit): models operand fetch for MAC/SIMD
  - WRITE (compute unit → SPMem): models result store from MAC/SIMD

Each channel has capacity=1: concurrent reads serialize, concurrent writes
serialize, but a read and a write can proceed in parallel.
"""
from __future__ import annotations

from dataclasses import dataclass

import simpy


@dataclass
class TcmRequest:
    """Request to read from or write to TCM."""

    direction: str      # "read" or "write"
    nbytes: int
    done: simpy.Event
    tag: str = ""       # optional label for debugging


class TcmBlock:
    """BW-serialized TCM model with dual read/write channels.

    Args:
        env: SimPy environment.
        read_bw_gbs: read bandwidth in GB/s (SPMem → compute).
        write_bw_gbs: write bandwidth in GB/s (compute → SPMem).
    """

    def __init__(
        self,
        env: simpy.Environment,
        read_bw_gbs: float = 512.0,
        write_bw_gbs: float = 512.0,
    ) -> None:
        self.env = env
        self._read_bw = read_bw_gbs
        self._write_bw = write_bw_gbs
        self._read_res = simpy.Resource(env, capacity=1)
        self._write_res = simpy.Resource(env, capacity=1)
        self._port: simpy.Store = simpy.Store(env)

    @property
    def port(self) -> simpy.Store:
        """The SimPy Store that blocks send TcmRequests to."""
        return self._port

    def _run(self):
        """Main process: receive TcmRequests, dispatch to channel processes."""
        while True:
            req: TcmRequest = yield self._port.get()
            if req is None:
                break
            self.env.process(self._handle(req))

    def _handle(self, req: TcmRequest):
        """Acquire channel, apply BW-based delay, signal done."""
        if req.direction == "write":
            res = self._write_res
            bw = self._write_bw
        else:
            res = self._read_res
            bw = self._read_bw

        with res.request() as lock:
            yield lock
            if bw > 0 and req.nbytes > 0:
                delay_ns = req.nbytes / bw
                yield self.env.timeout(delay_ns)
            req.done.succeed()
