import os
import sys
from pathlib import Path
import logging
import random

from dataclasses import dataclass

import cocotb
import cocotb.handle
from cocotb.queue import Queue
from cocotb.clock import Clock
from cocotb.triggers import Trigger, First, Event, Timer, RisingEdge, Edge, FallingEdge
from cocotb_tools.runner import get_runner

if cocotb.simulator.is_running():
    from spi_models import SpiInterface, SpiConfig, SpiSlave

ADDR_CTRL = 0
ADDR_INTMASK = 1
ADDR_STATUS = 2
ADDR_DATA = 3

INTMASK_RX_BUFFER_FULL = 0
INTMASK_TX_BUFFER_EMPTY = 1
INTMASK_LOST_RX_DATA = 3

STATUS_RX_BUFFER_FULL = 0
STATUS_TX_BUFFER_EMPTY = 1
STATUS_BUSY = 2
STATUS_LOST_RX_DATA = 3

CTRL_EN = 0
CTRL_MASTER = 1
CTRL_TX_EN = 2
CTRL_RX_EN = 3
CTRL_CLOCK_POLARITY = 4
CTRL_CLOCK_PHASE = 5
CTRL_PULSE_CSN = 6
CTRL_LSBFIRST = 7
CTRL_SIZE_SEL = 10
CTRL_DIV_SEL = 20

async def init(dut, master: int = 1, tx_en: int = 1):
    dut._log.info("Init started!")
    dut.waddress_i.value = 0
    dut.raddress_i.value = 0
    dut.wdata_i.value = 0
    dut.write_i.value = 0
    dut.read_i.value = 0
    dut.rst_in.value = 0

    await FallingEdge(dut.clk_i)
    await FallingEdge(dut.clk_i)

    # Release reset
    dut.rst_in.value = 1

    await FallingEdge(dut.clk_i)
    await FallingEdge(dut.clk_i)

    dut._log.info("Init done!")

class DutDriver:
    def __init__(self, dut):
        self.dut = dut
        self.irq_handlers = []

    async def clear_lost_rx_data(self):
        await self.write(ADDR_INTMASK, INTMASK_LOST_RX_DATA)

    async def configure_intmask(self, rx_buffer_full, tx_buffer_full, lost_rx_data):
        await self.write(ADDR_INTMASK,
                         (rx_buffer_full << INTMASK_RX_BUFFER_FULL) |
                         (tx_buffer_full << INTMASK_TX_BUFFER_FULL) |
                         (lost_rx_data << INTMASK_LOST_RX_DATA))

    async def configure(self, en, master, tx_en, rx_en, clock_polarity, clock_phase, pulse_csn, lsbfirst, size_sel, div_sel):
        await self.write(ADDR_CTRL,
                   (en << CTRL_EN) |
                   (master << CTRL_MASTER) |
                   (tx_en << CTRL_TX_EN) |
                   (rx_en << CTRL_RX_EN) |
                   (clock_polarity << CTRL_CLOCK_POLARITY) |
                   (clock_phase << CTRL_CLOCK_PHASE) |
                   (lsbfirst << CTRL_LSBFIRST) |
                   (size_sel << CTRL_SIZE_SEL) |
                   (div_sel << CTRL_DIV_SEL))

    async def read(self, address):
        await FallingEdge(self.dut.clk_i)
        self.dut.raddress_i.value = address
        self.dut.read_i.value = 1
        await FallingEdge(self.dut.clk_i)
        self.dut.read_i.value = 0
        return self.dut.rdata_o.value

    async def write(self, address, data):
        await FallingEdge(self.dut.clk_i)
        self.dut.waddress_i.value = address
        self.dut.wdata_i.value = data
        self.dut.write_i.value = 1
        await FallingEdge(self.dut.clk_i)
        self.dut.write_i.value = 0

    def register_interrupt_handler(self, handler):
        self.irq_handlers.push(handler)

    async def coroutine(self):
        await RisingEdge(self.dut.interrupt_o)

        if self.irq_handlers.len() == 0:
            raise Exception("Got interrupt but there is no irq handler. This would be a dead loop.")

        # The handlers are called until interrupt goes down!
        while self.dut.interrupt_o.value == 1:
            for irq_handler in irq_handlers:
                irq_handler(self)

@cocotb.test()
async def single_transission(dut):
    clk = Clock(dut.clk_i, 5, "ns")
    interface = SpiInterface(dut.csn_o, dut.sck_o, dut.miso_i, dut.mosi_o)
    config = SpiConfig(16, RisingEdge, FallingEdge, 40, "ns")
    slave = SpiSlave(interface, config)
    driver = DutDriver(dut)

    await cocotb.start(clk.start())

    await init(dut)

    await cocotb.start(slave.coroutine())
    await cocotb.start(driver.coroutine())

    await driver.configure(
        en = 1,
        master = 1,
        tx_en = 1,
        rx_en = 1,
        clock_polarity = 1,
        clock_phase = 1,
        pulse_csn = 0,
        lsbfirst = 0,
        size_sel = 1,
        div_sel = 2)

    # From slave point of view
    rx = random.randint(0, 255)
    tx = random.randint(0, 255)

    await slave.send_data(tx, 16)
    await slave.expect_transaction_in(100, "ns")

    await driver.write(ADDR_DATA, rx)

    await slave.wait_all()

    await FallingEdge(dut.clk_i)

    dut_received = await driver.read(ADDR_DATA)
    assert (int(dut_received) & 0xFF) == tx

    received = await slave.received_data()
    assert int(received) & 0xFF == rx

    await Timer(100, "ns")


def spi_peripheral_tests_runner():
    hdl_toplevel_lang = "vhdl"
    sim = os.getenv("SIM", "questa")

    proj_path = Path(__file__).resolve().parent.parent
    # equivalent to setting the PYTHONPATH environment variable
    sys.path.append(str(proj_path / "models"))

    sources = [
        proj_path / "src" / "spi_pkg.vhd",
        proj_path / "src" / "rs_latch.vhd",
        proj_path / "src" / "register.vhd",
        proj_path / "src" / "shift_register.vhd",
        proj_path / "src" / "spi_clkgen.vhd",
        proj_path / "src" / "spi_clkmon.vhd",
        proj_path / "src" / "spi_slave_ctrl.vhd",
        proj_path / "src" / "spi_master_ctrl.vhd",
        proj_path / "src" / "spi_master.vhd",
        proj_path / "src" / "spi_masterslave.vhd",
        proj_path / "src" / "spi_peripheral.vhd"
    ]

    build_args = []
    extra_args = []
    if sim == "ghdl":
        extra_args = ["--std=08"]
    elif sim == "questa" or sim == "modelsim":
        build_args = ["-2008"]

    # equivalent to setting the PYTHONPATH environment variable
    sys.path.append(str(proj_path / "tests"))

    runner = get_runner(sim)
    runner.build(
        sources=sources,
        hdl_toplevel="spi_peripheral",
        always=True,
        build_args=build_args + extra_args,
    )
    runner.test(
        hdl_toplevel="spi_peripheral", hdl_toplevel_lang=hdl_toplevel_lang,
        test_module="test_spi_peripheral", test_args = extra_args,
        gui = True
    )


if __name__ == "__main__":
    spi_peripheral_tests_runner()