"""IPCQ schemas and exceptions (ADR-0023 D2.5, D12, D14 F1).

This module contains the data structures and exceptions used by the
PE-level IPCQ collective communication infrastructure. The host-facing
sideband fan-out message ``IpcqInitMsg`` lives in
``kernbench.runtime_api.kernel`` (alongside other fabric messages),
while all internal token / metadata / command schemas are kept here.

Layering:
    PE_CPU       --IpcqRequest(IpcqSendCmd|IpcqRecvCmd)--> PE_IPCQ
    PE_IPCQ      --IpcqDmaToken-->                         PE_DMA (vc_comm)
    PE_DMA       --IpcqMetaArrival-->                      PE_IPCQ (atomic, D9)
    PE_IPCQ      --IpcqCreditMetadata-->                   peer PE_IPCQ (fast path, D9)

See ADR-0023 for the full design.
"""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Union

if TYPE_CHECKING:
    import simpy


# ── D14 F1: invalid direction exception ──────────────────────────────


class IpcqInvalidDirection(ValueError):
    """Raised when a kernel calls tl.send/recv with a direction that
    has no neighbor installed for this PE."""


# ── ADR-0023 D9.7: IPCQ slot-memory latency model ───────────────────
#
# Per-tier (bw_gbs, overhead_ns) used to charge the slot write (inbound)
# and slot read (recv consume). Mirrors topology.yaml component values.
_BUFFER_KIND_BW: dict[str, tuple[float, float]] = {
    "tcm":  (512.0, 0.0),
    "sram": (512.0, 2.0),
    "hbm":  (256.0, 6.0),
}


def slot_io_latency_ns(buffer_kind: str, nbytes: int) -> float:
    """Per-access latency for one slot read/write of ``nbytes`` against
    the IPCQ backing memory tier (``buffer_kind``)."""
    bw_gbs, overhead_ns = _BUFFER_KIND_BW.get(
        buffer_kind, _BUFFER_KIND_BW["tcm"],
    )
    return float(nbytes) / bw_gbs + overhead_ns


# ── D2.5: IpcqEndpoint ───────────────────────────────────────────────


@dataclass(frozen=True)
class IpcqEndpoint:
    """송신 측이 peer's rx_buffer 주소를 계산하기 위해 필요한 모든 정보 (D2.5).

    Sender PE_IPCQ uses this to compute the destination PA for its DMA
    write into the peer's rx ring buffer slot:

        slot_idx = sender.my_head % peer.n_slots
        dst_pa   = peer.rx_base_pa + slot_idx * peer.slot_size
    """

    sip: int                     # destination SIP
    cube: int                    # destination cube
    pe: int                      # destination PE (cube-local index)
    buffer_kind: str             # "tcm" | "hbm" | "sram"
    rx_base_pa: int              # peer rx_buffer base PA (PhysAddr.encode())
    rx_base_va: int              # peer rx_buffer base VA (optional, MMU)
    n_slots: int                 # peer ring depth (wrap-around modulo)
    slot_size: int               # peer slot size (offset multiplier)


# ── D12: IpcqInitEntry (used by IpcqInitMsg in kernel.py) ────────────


@dataclass(frozen=True)
class IpcqInitEntry:
    """One direction's neighbor entry that backend installs into a PE_IPCQ
    via IpcqInitMsg (kernbench.runtime_api.kernel.IpcqInitMsg, D12).
    """

    direction: str               # "N" | "S" | "E" | "W"
    peer: IpcqEndpoint           # see D2.5
    my_rx_base_pa: int           # this PE's own rx_buffer base
    my_rx_base_va: int           # this PE's own rx_buffer base VA (optional)
    n_slots: int                 # this PE's ring depth
    slot_size: int               # this PE's slot size
    # Credit fast path channel (D9).
    # Contract: must be a simpy.Store instance dedicated to receiving
    # IpcqCreditMetadata objects only. Backend wires it once at init time
    # and the receiving PE_IPCQ owns its consumer side; the sender (peer's
    # PE_IPCQ) puts IpcqCreditMetadata directly into this store via
    # _delayed_credit_send. Do not put any other object type.
    peer_credit_store: "simpy.Store"


# ── D12: IpcqSendCmd (PE_CPU → PE_IPCQ) ──────────────────────────────


@dataclass(frozen=True)
class IpcqSendCmd:
    """tl.send command issued by the kernel to PE_IPCQ."""

    direction: str               # "N" | "S" | "E" | "W"
    src_addr: int                # source data address (TCM/HBM/SRAM)
    src_space: str               # "tcm" | "hbm" | "sram"
    nbytes: int
    shape: tuple[int, ...]       # data shape (op_log + MemoryStore use)
    dtype: str
    handle_id: str               # completion tracking
    # In-flight data snapshot captured at tl.send() time from the
    # TensorHandle.data field. Carries the actual numpy array that was
    # visible at recv-time (when handle.data was populated), avoiding a
    # Phase 1 race where a later IPCQ inbound overwrites the sender's
    # slot between recv and send. If None, PE_DMA outbound falls back to
    # reading MemoryStore[src_addr] (correct for sources that are never
    # overwritten, such as HBM tiles).
    data: Any = None
    data_op: bool = True         # ADR-0020 op_log recording flag

    @property
    def logical_bytes(self) -> int:
        # framing + direction enum + src_addr + src_space enum + nbytes
        # + shape (len marker + 4·rank) + dtype tag (ADR-0064 D2).
        return 4 + 1 + 4 + 1 + 4 + (1 + 4 * len(self.shape)) + 1


# ── D12: IpcqRecvCmd (PE_CPU → PE_IPCQ) ──────────────────────────────


@dataclass(frozen=True)
class IpcqRecvCmd:
    """tl.recv command issued by the kernel to PE_IPCQ.

    Two modes (recv_mode):
        "return_slot" — return slot address as-is (default, zero-copy).
                        Kernel uses the slot memory directly.
        "copy_to_dst" — copy slot data to dst_addr, then return.

    ``consume`` (DIAGNOSTIC ONLY): when False, recv still blocks until the
    payload lands in the slot, but skips the slot-read latency charge
    (slot-IO + PE↔bank fabric drain for SRAM/HBM tiers). This exists
    solely so the pe2pe overview plot can compare apples-to-apples
    against tl.store (a one-sided write that pays no read on DST). Real
    kernels always need the data they receive — leave this True.
    """

    direction: str | None        # None → round-robin (weak fairness, D4)
    shape: tuple[int, ...]
    dtype: str
    handle_id: str
    recv_mode: str = "return_slot"
    dst_addr: int = 0            # used only when recv_mode == "copy_to_dst"
    dst_space: str = ""          # used only when recv_mode == "copy_to_dst"
    blocking: bool = True
    data_op: bool = True
    consume: bool = True         # DIAGNOSTIC: see docstring

    @property
    def logical_bytes(self) -> int:
        # framing + direction enum + shape (len marker + 4·rank) + dtype tag
        # + dst_addr + dst_space enum (ADR-0064 D2).
        return 4 + 1 + (1 + 4 * len(self.shape)) + 1 + 4 + 1


# ── D12: IpcqDmaToken (PE_IPCQ → PE_DMA, vc_comm) ───────────────────


@dataclass
class IpcqDmaToken:
    """Token sent from PE_IPCQ to PE_DMA (vc_comm channel) carrying both
    the data move request and the piggyback metadata (ADR-0023 D9).

    Receiving PE_DMA processes this atomically (I6 MUST):
        1. MemoryStore.write(dst_endpoint.buffer_kind, dst_addr, data)
        2. Forward IpcqMetaArrival(token=self) to peer PE_IPCQ
    No yield is allowed between the two steps.

    The ``data`` field is a snapshot taken by the sender's PE_DMA at the
    moment the send is issued. This preserves "in-flight data" semantics:
    if the sender mutates its source memory after issuing the send but
    before arrival, the receiver still gets the snapshot. The snapshot is
    None for control-only tokens (e.g. credit-only updates).
    """

    # ── Data movement (single-hop DMA write) ──
    src_addr: int
    src_space: str
    dst_addr: int                # already-computed peer rx slot PA
    dst_endpoint: IpcqEndpoint   # routing target (sip/cube/pe) + buffer_kind
    nbytes: int
    handle_id: str               # completion notify back to sender PE_IPCQ
    # Optional shape/dtype carried for op_log + MemoryStore convenience.
    shape: tuple[int, ...] = ()
    dtype: str = "f16"
    # In-flight data snapshot (sender PE_DMA captures this at send time).
    data: Any = None

    # ── Piggyback metadata (D9) ──
    sender_seq: int = 0          # monotonic; receiver updates peer_head_cache
    src_sip: int = 0
    src_cube: int = 0
    src_pe: int = 0
    src_direction: str = "E"     # sender-side direction; receiver maps to its own

    data_op: bool = True


# ── D12: IpcqMetaArrival (PE_DMA → PE_IPCQ, intra-PE wire) ──────────


@dataclass
class IpcqMetaArrival:
    """Posted by receiving PE_DMA into the destination PE's PE_IPCQ inbox
    in the same SimPy step as the MemoryStore.write (D9, I6 MUST).

    The receiver PE_IPCQ uses ``token.sender_seq`` to update its
    peer_head_cache for the corresponding direction.
    """

    token: IpcqDmaToken


# ── D12: IpcqCreditMetadata (PE_IPCQ → peer PE_IPCQ, fast path) ─────


@dataclass(frozen=True)
class IpcqCreditMetadata:
    """Credit return — recv-side → send-side fast path (D9).

    Sent by ``PeIpcqComponent._delayed_credit_send`` after a
    bottleneck-BW based latency, putting the metadata directly into
    the peer's pre-wired credit store (no fabric routing).

    ``dst_rx_base_pa`` is the receiver's ``my_rx_base_pa`` for the direction
    whose slot was consumed. The original sender matches this against
    ``qp.peer.rx_base_pa`` to find the correct direction (ADR-0025 D3) —
    unambiguous even when multiple directions share the same peer (e.g.
    2-rank bidirectional ring).
    """

    consumer_seq: int            # my_tail at recv side (new tail value)
    dst_rx_base_pa: int          # receiver-side my_rx_base_pa (ADR-0025 D3)
    src_sip: int                 # which peer is sending the credit (diag)
    src_cube: int
    src_pe: int
    src_direction: str           # sender-side direction (peer maps to its own)


# ── Request wrapper (PE_CPU → PE_IPCQ) ───────────────────────────────


@dataclass
class IpcqRequest:
    """Wrapper carrying an IpcqSendCmd or IpcqRecvCmd plus a SimPy completion
    event. Posted by PE_CPU into PE_IPCQ's inbox; PE_IPCQ calls
    ``done.succeed()`` when the request is fully processed.

    For recv requests, the result (slot address, direction, dtype, shape)
    is written into ``result_data`` so the caller can read it after wait.
    """

    command: "IpcqSendCmd | IpcqRecvCmd"
    done: "simpy.Event"
    result_data: dict[str, Any] = field(default_factory=dict)


# ── RecvFuture (kernel ↔ runner handshake for tl.recv_async / tl.wait) ─


@dataclass
class RecvFuture:
    """Opaque future returned by ``tl.recv_async``.

    The KernelRunner attaches a SimPy event and the IpcqRequest in the
    background; ``tl.wait(future)`` switches back to the runner which
    yields on the event and resolves the result into a TensorHandle.
    """

    cmd: "IpcqRecvCmd"
    request: Any = None         # IpcqRequest (set by runner)
    event: Any = None           # simpy.Event (set by runner)
    resolved: bool = False
    result: Any = None          # cached TensorHandle after wait()
