"""Data types for pe_accel_v1: descriptors, triggers, tile commands.

All types are frozen/plain dataclasses with no logic.
Schedule generators live in tiling/schedule.py.
"""
from __future__ import annotations

from dataclasses import dataclass
from enum import Enum, auto


# -- Enums ---------------------------------------------------------------------

class CmdType(Enum):
    DMA_LOAD = auto()
    TENSOR_OP = auto()
    DMA_FLUSH = auto()


# -- Inter-block messaging -----------------------------------------------------

@dataclass
class Trigger:
    """Completion token passed between hardware blocks."""

    tile_id: int
    pipeline_id: int
    vc: int | None = None
    source_block: str = ""


@dataclass
class DmaRequest:
    """DMA load request — descriptor lookup key only.

    Transfer params (size, address) live in the pre-loaded DmaInDescriptor.
    """

    tile_id: int
    pipeline_id: int
    operand: str  # "A" or "B"


# -- Descriptors (pre-loaded by pipelines, consumed by blocks) -----------------

@dataclass
class DmaInDescriptor:
    """Per-operand DMA read descriptor."""

    size_bytes: int
    src_addr: int = 0
    next_block: str = "GEMM"   # "GEMM" | "MATH" | "COMPLETION"


@dataclass
class GemmDescriptor:
    """Per-tile GEMM descriptor."""

    Tm: int
    Tk: int
    Tn: int
    triggers_needed: int = 2   # 2 = both operands from DMA; 1 = SPMem bypass
    gemm_load: bool = True
    gemm_compute: bool = True
    next_block: str = "MATH"   # "MATH" | "DMAWB" | "DONE"


@dataclass
class MathDescriptor:
    """Per-tile K-accumulation descriptor (used by GEMM pipeline)."""

    Tm: int
    Tn: int
    is_last_k: bool
    skip_dmaWB: bool   # True = C stays in SPMem; False = flush to HBM


@dataclass
class MathOpDescriptor:
    """Per-tile element-wise math op descriptor (used by math pipeline)."""

    Tm: int
    Tn: int
    op: str            # "exp", "log", "sqrt", "sigmoid", etc.
    src_addr: int = 0
    dst_addr: int = 0


@dataclass
class DmaWBDescriptor:
    """Per-tile DMA writeback descriptor."""

    Tm: int
    Tn: int
    dst_addr: int = 0


# -- Tile commands (produced by schedule generators) ---------------------------

@dataclass
class TileCommand:
    """A single GEMM tile command."""

    cmd_id: int
    cmd_type: CmdType
    tile_id: int
    m_idx: int
    k_idx: int
    n_idx: int
    is_last_k: bool = False
    pipeline_id: int = 0
    a_tile_addr: int = 0
    b_tile_addr: int = 0
    c_tile_addr: int = 0


@dataclass
class SchedulePlan:
    """Full tile schedule for one GEMM operation."""

    commands: list  # list[TileCommand]
    M_tiles: int
    K_tiles: int
    N_tiles: int


@dataclass
class MathTileCommand:
    """A single element-wise math tile command."""

    cmd_id: int
    tile_id: int
    m_idx: int
    n_idx: int
    src_tile_addr: int = 0
    dst_tile_addr: int = 0
    pipeline_id: int = 0


@dataclass
class MathSchedulePlan:
    """Full tile schedule for one element-wise math operation."""

    commands: list  # list[MathTileCommand]
    M_tiles: int
    N_tiles: int


