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_STATUS = 1
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
REG_ERROR = 8
REG_DESC_READ_ADDR = 9
REG_DESC_READ_LEN = 10
REG_DESC_READ_CONFIG = 11
REG_READ_STATUS = 12
REG_DESC_WRITE_ADDR = 13
REG_DESC_WRITE_LEN_AND_TAG = 14
REG_STATUS_WRITE_LEN = 15
REG_STATUS_WRITE_CONFIG = 16


# REG_CONTROL pulse bits
CTRL_START = 0
CTRL_RST_SOFT = 1
CTRL_CFG_BUS_VALID = 2
CTRL_SEND_DESC_READ = 3
CTRL_SEND_DESC_WRITE = 4
CTRL_TAKE_STATUS_READ = 5
CTRL_TAKE_STATUS_WRITE = 6


def reg_addr(reg_index: int) -> int:
    # AXI-Lite uses byte addresses, 32-bit registers are spaced by 4 bytes.
    return reg_index * 4


def u32(value: int) -> bytes:
    return int(value & 0xFFFFFFFF).to_bytes(4, "little")


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)

    async def write_reg(self, reg_index: int, value: int):
        await self.axil.write(reg_addr(reg_index), u32(value))

    async def read_reg(self, reg_index: int) -> int:
        resp = await self.axil.read(reg_addr(reg_index), 4)
        return int.from_bytes(bytes(resp.data), "little")


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

    await tb.write_reg(REG_DAC_WIDTH, 0x0000_0123)

    value = await tb.read_reg(REG_DAC_WIDTH)
    assert value == 0x0000_0123


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

    await tb.write_reg(REG_DAC_WIDTH, 0x10)
    await tb.write_reg(REG_DAC_PERIOD, 0x40)
    await tb.write_reg(REG_DAC_PULSE_NUM, 3)
    await tb.write_reg(REG_DAC_PULSE_HEIGHT, 0x7FF)
    await tb.write_reg(REG_ADC_PERIOD, 0x80)
    await tb.write_reg(REG_WINDOW_SIZE, 16)

    assert await tb.read_reg(REG_DAC_WIDTH) == 0x10
    assert await tb.read_reg(REG_WINDOW_SIZE) == 16

    # W1S pulse: bit 2 = cfg_bus_valid_o.
    await tb.write_reg(REG_CONTROL, 1 << CTRL_CFG_BUS_VALID)

    await RisingEdge(dut.ctrl_clk)


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

    # mock
    await tb.write_reg(REG_DESC_READ_ADDR, 0x0000_1000)
    await tb.write_reg(REG_DESC_READ_LEN, 64)
    await tb.write_reg(REG_DESC_READ_CONFIG, 0x0000_0001)

    assert await tb.read_reg(REG_DESC_READ_ADDR) == 0x0000_1000
    assert await tb.read_reg(REG_DESC_READ_LEN) == 64

    # W1S pulse: bit 3 = send_desc_read_o.
    await tb.write_reg(REG_CONTROL, 1 << CTRL_SEND_DESC_READ)

    for _ in range(3):
        await RisingEdge(dut.ctrl_clk)