"""Tests for IPCQ type schemas (ADR-0023 D2.5, D12, D14 F1)."""
import pytest

from kernbench.common.ipcq_types import (
    IpcqCreditMetadata,
    IpcqDmaToken,
    IpcqEndpoint,
    IpcqInitEntry,
    IpcqInvalidDirection,
    IpcqMetaArrival,
    IpcqRecvCmd,
    IpcqSendCmd,
)
from kernbench.runtime_api.kernel import IpcqInitMsg


# ── IpcqEndpoint ─────────────────────────────────────────────────────


def test_ipcq_endpoint_basic():
    ep = IpcqEndpoint(
        sip=0, cube=0, pe=1,
        buffer_kind="tcm",
        rx_base_pa=0x1000, rx_base_va=0,
        n_slots=8, slot_size=4096,
    )
    assert ep.sip == 0
    assert ep.buffer_kind == "tcm"
    assert ep.n_slots == 8


def test_ipcq_endpoint_frozen():
    ep = IpcqEndpoint(
        sip=0, cube=0, pe=1, buffer_kind="tcm",
        rx_base_pa=0x1000, rx_base_va=0, n_slots=8, slot_size=4096,
    )
    with pytest.raises(Exception):  # FrozenInstanceError
        ep.sip = 99  # type: ignore


# ── IpcqDmaToken ─────────────────────────────────────────────────────


def test_ipcq_dma_token():
    ep = IpcqEndpoint(
        sip=0, cube=0, pe=1, buffer_kind="tcm",
        rx_base_pa=0x1000, rx_base_va=0, n_slots=8, slot_size=4096,
    )
    tok = IpcqDmaToken(
        src_addr=0x500, src_space="tcm",
        dst_addr=0x1000, dst_endpoint=ep,
        nbytes=128, handle_id="h1",
        sender_seq=0,
        src_sip=0, src_cube=0, src_pe=0, src_direction="E",
    )
    assert tok.nbytes == 128
    assert tok.dst_endpoint.buffer_kind == "tcm"
    assert tok.data_op is True


# ── IpcqCreditMetadata ───────────────────────────────────────────────


def test_ipcq_credit_metadata():
    cm = IpcqCreditMetadata(
        consumer_seq=3, dst_rx_base_pa=0x1000,
        src_sip=0, src_cube=0, src_pe=1, src_direction="W",
    )
    assert cm.consumer_seq == 3
    assert cm.src_direction == "W"


def test_ipcq_credit_metadata_frozen():
    cm = IpcqCreditMetadata(
        consumer_seq=3, dst_rx_base_pa=0x1000,
        src_sip=0, src_cube=0, src_pe=1, src_direction="W",
    )
    with pytest.raises(Exception):
        cm.consumer_seq = 99  # type: ignore


# ── IpcqMetaArrival ──────────────────────────────────────────────────


def test_ipcq_meta_arrival():
    ep = IpcqEndpoint(
        sip=0, cube=0, pe=1, buffer_kind="tcm",
        rx_base_pa=0x1000, rx_base_va=0, n_slots=8, slot_size=4096,
    )
    tok = IpcqDmaToken(
        src_addr=0x500, src_space="tcm",
        dst_addr=0x1000, dst_endpoint=ep,
        nbytes=128, handle_id="h1",
        sender_seq=0,
        src_sip=0, src_cube=0, src_pe=0, src_direction="E",
    )
    ma = IpcqMetaArrival(token=tok)
    assert ma.token.sender_seq == 0
    assert ma.token.src_direction == "E"


# ── IpcqSendCmd / IpcqRecvCmd ────────────────────────────────────────


def test_ipcq_send_cmd():
    cmd = IpcqSendCmd(
        direction="E", src_addr=0x100, src_space="tcm",
        nbytes=64, shape=(8, 8), dtype="f16", handle_id="s1",
    )
    assert cmd.direction == "E"
    assert cmd.data_op is True


def test_ipcq_recv_cmd_default_return_slot():
    cmd = IpcqRecvCmd(direction="W", shape=(8, 8), dtype="f16", handle_id="r1")
    assert cmd.recv_mode == "return_slot"
    assert cmd.dst_addr == 0


def test_ipcq_recv_cmd_round_robin():
    cmd = IpcqRecvCmd(direction=None, shape=(8, 8), dtype="f16", handle_id="r2")
    assert cmd.direction is None


def test_ipcq_recv_cmd_copy_to_dst():
    cmd = IpcqRecvCmd(
        direction="W", recv_mode="copy_to_dst",
        dst_addr=0x2000, dst_space="hbm",
        shape=(8, 8), dtype="f16", handle_id="r3",
    )
    assert cmd.recv_mode == "copy_to_dst"
    assert cmd.dst_addr == 0x2000


# ── IpcqInvalidDirection ─────────────────────────────────────────────


def test_ipcq_invalid_direction():
    with pytest.raises(IpcqInvalidDirection):
        raise IpcqInvalidDirection("direction 'X' not installed")


# ── IpcqInitEntry / IpcqInitMsg ──────────────────────────────────────


def test_ipcq_init_entry_and_msg():
    import simpy
    env = simpy.Environment()
    credit_store = simpy.Store(env)

    ep = IpcqEndpoint(
        sip=0, cube=0, pe=1, buffer_kind="tcm",
        rx_base_pa=0x1000, rx_base_va=0, n_slots=8, slot_size=4096,
    )
    entry = IpcqInitEntry(
        direction="E", peer=ep,
        my_rx_base_pa=0x2000, my_rx_base_va=0,
        n_slots=8, slot_size=4096,
        peer_credit_store=credit_store,
    )
    msg = IpcqInitMsg(
        correlation_id="c1", request_id="r1",
        target_sips=(0,), target_cubes=(0,), target_pe=0,
        entries=(entry,),
        backpressure_mode="sleep",
        buffer_kind="tcm",
        credit_size_bytes=16,
    )
    assert msg.entries[0].direction == "E"
    assert msg.entries[0].peer.sip == 0
    assert msg.credit_size_bytes == 16
