import cocotb
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge
from cocotbext.axi import AxiLiteBus, AxiLiteMaster


# register indexes from axi4l_reg_map_controller_pkg.sv
REG_CONTROL = 0
REG_DAC_WIDTH = 2
REG_DAC_PERIOD = 3
REG_DAC_PULSE_NUM = 4
REG_DAC_PULSE_HEIGHT = 5
REG_ADC_PERIOD = 6
REG_WINDOW_SIZE = 7


def reg_addr(reg_index: int) -> int:
    # AXI-Lite usually uses byte addresses for 32-bit regs
    return reg_index * 4


class TB:
    def __init__(self, dut):
        self.dut = dut

        cocotb.start_soon(Clock(dut.ctrl_clk, 10, units="ns").start())

        self.axil = AxiLiteMaster(
            AxiLiteBus.from_prefix(dut, "s_axil"),
            dut.ctrl_clk,
            dut.rst,
        )

    async def reset(self):
        self.dut.rst.value = 1
        for _ in range(5):
            await RisingEdge(self.dut.ctrl_clk)

        self.dut.rst.value = 0
        for _ in range(5):
            await RisingEdge(self.dut.ctrl_clk)


@cocotb.test()
async def simple_axil_write(dut):
    tb = TB(dut)
    await tb.reset()

    # Простая запись в RW-регистр REG_DAC_WIDTH.
    await tb.axil.write(reg_addr(REG_DAC_WIDTH), (0x0000_0123).to_bytes(4, "little"))

    # Не обязательно, но удобно проверить, что запись прошла через AXI-Lite.
    resp = await tb.axil.read(reg_addr(REG_DAC_WIDTH), 4)
    value = int.from_bytes(bytes(resp.data), "little")
    assert value == 0x0000_0123

    # Пример W1S-записи в control: bit 3 = cfg_bus_valid pulse.
    # Дальше уже можно проверять твою логику контроллера.
    await tb.axil.write(reg_addr(REG_CONTROL), (1 << 3).to_bytes(4, "little"))

    await RisingEdge(dut.ctrl_clk)