import logging
import cocotb

from cocotb.triggers import Timer, FallingEdge
from cocotb.clock import Clock

from cocotbext.axi import AxiStreamBus
from cocotbext.pcie.core import RootComplex
from cocotbext.pcie.xilinx.us import UltraScalePlusPcieDevice

CLK_PERIOD = 4

class TB:

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

        self.log = logging.getLogger("cocotb.tb")
        self.log.setLevel(logging.INFO)

        self.rc = RootComplex()

        self.dev = UltraScalePlusPcieDevice(
            # configuration options
            pcie_generation=3,
            # pcie_link_width=2,
            # user_clk_frequency=250e6,
            alignment="dword",
            cq_straddle=False,
            cc_straddle=False,
            rq_straddle=False,
            rc_straddle=False,
            rc_4tlp_straddle=False,
            pf_count=1,
            max_payload_size=1024,
            enable_client_tag=True,
            enable_extended_tag=True,
            enable_parity=False,
            enable_rx_msg_interface=False,
            enable_sriov=False,
            enable_extended_configuration=False,

            pf0_msi_enable=True,
            pf0_msi_count=32,
            pf1_msi_enable=False,
            pf1_msi_count=1,
            pf2_msi_enable=False,
            pf2_msi_count=1,
            pf3_msi_enable=False,
            pf3_msi_count=1,
            pf0_msix_enable=False,
            pf0_msix_table_size=0,
            pf0_msix_table_bir=0,
            pf0_msix_table_offset=0x00000000,
            pf0_msix_pba_bir=0,
            pf0_msix_pba_offset=0x00000000,
            pf1_msix_enable=False,
            pf1_msix_table_size=0,
            pf1_msix_table_bir=0,
            pf1_msix_table_offset=0x00000000,
            pf1_msix_pba_bir=0,
            pf1_msix_pba_offset=0x00000000,
            pf2_msix_enable=False,
            pf2_msix_table_size=0,
            pf2_msix_table_bir=0,
            pf2_msix_table_offset=0x00000000,
            pf2_msix_pba_bir=0,
            pf2_msix_pba_offset=0x00000000,
            pf3_msix_enable=False,
            pf3_msix_table_size=0,
            pf3_msix_table_bir=0,
            pf3_msix_table_offset=0x00000000,
            pf3_msix_pba_bir=0,
            pf3_msix_pba_offset=0x00000000,

            # signals
            user_clk=dut.clk_250,
            user_reset=dut.rst_250,
            user_lnk_up=dut.user_lnk_up,

            rq_bus=AxiStreamBus.from_entity(dut.s_axis_rq),
            rc_bus=AxiStreamBus.from_entity(dut.m_axis_rc),
            cq_bus=AxiStreamBus.from_entity(dut.m_axis_cq),
            cc_bus=AxiStreamBus.from_entity(dut.s_axis_cc),
        )

        self.dev.functions[0].configure_bar(0, 64*1024)

        self.rc.make_port().connect(self.dev)

@cocotb.test
async def test_sanity(dut):
    tb = TB(dut)

    await FallingEdge(dut.rst_250)
    await Timer(100, 'ns')

    await tb.rc.enumerate()

    mem = tb.rc.mem_pool.alloc_region(16*1024*1024)

    dev = tb.rc.find_device(tb.dev.functions[0].pcie_id)
    await dev.enable_device()
    await dev.set_master()

    dev_bar0 = dev.bar_window[0]

    tb.log.info(dev_bar0.write)

    message = b"Hello, world! This is a long string of data with many letters and words."

    await mem.write(0, message)

    await dev_bar0.write_dword(0x0, 0x00000000)
    await dev_bar0.write_dword(0x4, 0x00000000)
    await dev_bar0.write_dword(0x8, 0x00000000)
    await dev_bar0.write_dword(0xc, len(message))
    await dev_bar0.write_dword(0x10, 0x00000001)

    await Timer(100, "ns")

    await dev_bar0.write_dword(0x20, 0x00000100)
    await dev_bar0.write_dword(0x24, 0x00000000)
    await dev_bar0.write_dword(0x28, 0x00000000)
    await dev_bar0.write_dword(0x2c, len(message))
    await dev_bar0.write_dword(0x30, 0x00000001)

    await Timer(10, "us")

    read_msg = await mem.read(0x100, len(message))

    tb.log.info(read_msg)