import cocotb
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge
from cocotbext.axi import AxiLiteBus, AxiLiteMaster
from cocotbext.axi import AxiStreamBus, AxiStreamBus, AxiStreamSource, AxiStreamSink, AxiStreamFrame
from cocotbext.axi import AxiBus, AxiRam

from reg_map import *

# Helpers

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")

# Drivers

class Drivers:
    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
        )

        self.axis_source = AxiStreamSource(
            AxiStreamBus.from_prefix(dut, "s_axis_write_data"),
            dut.ctrl_clk,
            dut.rst
        )

        self.axis_sink = AxiStreamSink(
            AxiStreamBus.from_prefix(dut, "m_axis_read_data"),
            dut.ctrl_clk,
            dut.rst
        )

        self.memory = AxiRam(
            AxiBus.from_prefix(dut, "m_axi"),
            dut.ctrl_clk,
            dut.rst
        )

    # Reset

    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)

    # AxiLite Driver

    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")
    
    async def pulse_control(self, mask):
        await self.write_reg( REG_CONTROL, mask )

    # Reflectometer Driver

    async def configure_reflectometer (
        self,
        pulse_width,
        pulse_period,
        pulse_num,
        pulse_height,
        adc_period,
        window_size
    ):
        await self.write_reg(REG_DAC_WIDTH, pulse_width)
        await self.write_reg(REG_DAC_PERIOD, pulse_period)
        await self.write_reg(REG_DAC_PULSE_NUM, pulse_num)
        await self.write_reg(REG_DAC_PULSE_HEIGHT, pulse_height)
        await self.write_reg(REG_ADC_PERIOD, adc_period)
        await self.write_reg(REG_WINDOW_SIZE, window_size)

        await self.pulse_control(CTRL_CFG_BUS_VALID)

    async def send_start(self):
        await self.pulse_control(CTRL_START)

    async def soft_reset(self):
        await self.pulse_control(CTRL_RST_SOFT)

    async def get_status(self):
        status = await self.read_reg(REG_STATUS)
        return {
        "busy": bool(status & STATUS_BUSY),
        "processing_done": bool(status & STATUS_PROCESSING_DONE),
        "desc_read_busy": bool(status & STATUS_DESC_READ_BUSY),
        "desc_write_busy": bool(status & STATUS_DESC_WRITE_BUSY),
        "status_read_busy": bool(status & STATUS_STATUS_READ_BUSY),
        "status_write_busy": bool(status & STATUS_STATUS_WRITE_BUSY),
        "desc_read_hs": bool(status & STATUS_DESC_READ_HS),
        "desc_write_hs": bool(status & STATUS_DESC_WRITE_HS),
        "status_read_hs": bool(status & STATUS_STATUS_READ_HS),
        "status_write_hs": bool(status & STATUS_STATUS_WRITE_HS)
        }

    async def wait_status(self, field, value=True):

        while True:

            status = await self.get_status()

            if status[field] == value:
                return

            await RisingEdge(self.dut.ctrl_clk)

    async def wait_processing_done(self):
        await self.wait_status("processing_done")


    async def wait_finish(self):
        await self.wait_status("busy", False)

    # DMA Driver

    async def send_desc_write(self, addr, length_tag):

        await self.write_reg(REG_DESC_WRITE_ADDR, addr)
        await self.write_reg(REG_DESC_WRITE_LEN_AND_TAG, length_tag)
        await self.pulse_control(CTRL_SEND_DESC_WRITE)

    async def send_desc_read(self, addr, length, config):

        await self.write_reg(REG_DESC_READ_ADDR, addr)
        await self.write_reg(REG_DESC_READ_LEN, length)
        await self.write_reg(REG_DESC_READ_CONFIG, config)
        await self.pulse_control( CTRL_SEND_DESC_READ)

    async def take_status_write(self):
        await self.pulse_control(CTRL_TAKE_STATUS_WRITE)

        return (
            await self.read_reg(REG_STATUS_WRITE_LEN),
            await self.read_reg(REG_STATUS_WRITE_CONFIG) )
    
    async def take_read_status(self):

        await self.pulse_control( CTRL_TAKE_STATUS_READ)
        return await self.read_reg(REG_READ_STATUS)
    
    async def wait_dma_write_done(self):
        await self.wait_status("desc_write_busy", False)

    async def wait_dma_read_done(self):
        await self.wait_status("desc_read_busy", False)

    async def wait_status_read_handshake(self):
        await self.wait_status("status_read_hs")

    async def wait_status_write_handshake(self):
        await self.wait_status("status_write_hs")

    # AxiStream Driver

    async def send_axis_data(self, data: bytes):

        await self.axis_source.send(AxiStreamFrame(data) )

    async def receive_axis_data(self):

        frame = await self.axis_sink.recv()
        return bytes(frame)

    # Memory Driver

    async def read_memory(self, addr, length):

        return await self.memory.read(addr, length )

    async def write_memory(self, addr, data ):

        await self.memory.write( addr, data )

    async def clear_memory( self, addr, length ):

        await self.memory.write( addr, bytes(length) )