"""GEMM Tiling Pipeline: splits (M,K)×(K,N) into tiles and coordinates execution.

Flow per tile:
  DMA_IN(A tile) + DMA_IN(B tile) → GEMM(fetch + compute) → MATH(K-accum) → DMA_WB

The pipeline does NOT own hardware blocks — it uses the component's shared
blocks via descriptor tables and SimPy queues.

Constructor starts two SimPy processes:
  - _feed_commands(): sends DmaRequests to shared dmaIN_cmd_q
  - _collect_completions(): waits for all output tiles to flush
"""
from __future__ import annotations

from math import ceil

import simpy

from kernbench.components.custom.pe_accel.scheduler.tiling import generate_gemm_tiles
from kernbench.components.custom.pe_accel.types import (
    CmdType,
    DmaInDescriptor,
    DmaRequest,
    DmaWBDescriptor,
    GemmDescriptor,
    MathDescriptor,
    Trigger,
)


class GemmPipeline:
    """Coordinates one tiled GEMM operation across shared hardware blocks."""

    def __init__(
        self,
        env: simpy.Environment,
        M: int, K: int, N: int,
        tile_m: int, tile_k: int, tile_n: int,
        bytes_per_element: int,
        pipeline_id: int,
        reply_queue: simpy.Store,
        dmaIN_cmd_q: simpy.Store,
        dmaIN_to_fetch_trig: simpy.Store,
        A_addr: int = 0,
        B_addr: int = 0,
        C_addr: int = 0,
        dma_a: bool = True,
        dma_b: bool = True,
        dma_c: bool = True,
    ) -> None:
        self.env = env
        self.M, self.K, self.N = M, K, N
        self.pipeline_id = pipeline_id
        self.reply_queue = reply_queue
        self.dmaIN_cmd_q = dmaIN_cmd_q
        self.dmaIN_to_fetch_trig = dmaIN_to_fetch_trig

        self._dma_a = dma_a
        self._dma_b = dma_b
        self._skip_dmaWB = not dma_c

        _Tm = min(tile_m, M)
        _Tk = min(tile_k, K)
        _Tn = min(tile_n, N)

        self.M_tiles = ceil(M / tile_m)
        self.K_tiles = ceil(K / tile_k)
        self.N_tiles = ceil(N / tile_n)

        triggers_per_tile = 2 if (dma_a and dma_b) else 1

        # Generate tile schedule with pre-computed addresses
        self.schedule = generate_gemm_tiles(
            self.M_tiles, self.K_tiles, self.N_tiles,
            M=M, K=K, N=N,
            tile_m=_Tm, tile_k=_Tk, tile_n=_Tn,
            bytes_per_element=bytes_per_element,
            A_addr=A_addr, B_addr=B_addr, C_addr=C_addr,
            pipeline_id=pipeline_id,
        )

        # Build descriptor tables for shared blocks
        pid = pipeline_id
        a_tile_bytes = _Tm * _Tk * bytes_per_element
        b_tile_bytes = _Tk * _Tn * bytes_per_element

        self.dmaIN_descs: dict[tuple, DmaInDescriptor] = {}
        self.gemm_descs: dict[tuple, GemmDescriptor] = {}
        self.math_descs: dict[tuple, MathDescriptor] = {}
        self.dmaWB_descs: dict[tuple, DmaWBDescriptor] = {}

        for cmd in self.schedule.commands:
            if cmd.cmd_type != CmdType.DMA_LOAD:
                continue
            t = cmd.tile_id

            if dma_a:
                self.dmaIN_descs[(pid, t, "A")] = DmaInDescriptor(
                    size_bytes=a_tile_bytes, src_addr=cmd.a_tile_addr
                )
            if dma_b:
                self.dmaIN_descs[(pid, t, "B")] = DmaInDescriptor(
                    size_bytes=b_tile_bytes, src_addr=cmd.b_tile_addr
                )

            self.gemm_descs[(pid, t)] = GemmDescriptor(
                Tm=_Tm, Tk=_Tk, Tn=_Tn,
                triggers_needed=triggers_per_tile,
                next_block="MATH",
            )

            self.math_descs[(pid, t)] = MathDescriptor(
                Tm=_Tm, Tn=_Tn,
                is_last_k=cmd.is_last_k,
                skip_dmaWB=self._skip_dmaWB,
            )
            if not self._skip_dmaWB and cmd.is_last_k:
                self.dmaWB_descs[(pid, t)] = DmaWBDescriptor(
                    Tm=_Tm, Tn=_Tn, dst_addr=cmd.c_tile_addr
                )

        self.expected_flushes = self.M_tiles * self.N_tiles
        self.completed_flushes = 0
        self.done_at: int = 0
        self.done: simpy.Event = env.event()

        env.process(self._feed_commands())
        env.process(self._collect_completions())

    def _feed_commands(self):
        """Send DmaRequests for each tile's operands to dmaIN_cmd_q."""
        for cmd in self.schedule.commands:
            if cmd.cmd_type != CmdType.DMA_LOAD:
                continue

            if self._dma_a:
                yield self.dmaIN_cmd_q.put(DmaRequest(
                    tile_id=cmd.tile_id, pipeline_id=self.pipeline_id, operand="A",
                ))
            if self._dma_b:
                yield self.dmaIN_cmd_q.put(DmaRequest(
                    tile_id=cmd.tile_id, pipeline_id=self.pipeline_id, operand="B",
                ))

            if not self._dma_a and not self._dma_b:
                yield self.dmaIN_to_fetch_trig.put(Trigger(
                    tile_id=cmd.tile_id, pipeline_id=self.pipeline_id, source_block="PIPELINE",
                ))

    def _collect_completions(self):
        """Wait for all output tile flush completions, then signal done."""
        while self.completed_flushes < self.expected_flushes:
            yield self.reply_queue.get()
            self.completed_flushes += 1

        self.done_at = int(self.env.now)
        self.done.succeed()
