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 = 1
INTMASK_TX_BUFFER_EMPTY = 2
INTMASK_LOST_RX_DATA = 4

STATUS_RX_BUFFER_FULL = 0
STATUS_TX_BUFFER_EMPTY = 2
STATUS_BUSY = 8
STATUS_LOST_RX_DATA = 16

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 = []
        self._handle_interrupt = False

    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) |
                   (pulse_csn << CTRL_PULSE_CSN) |
                   (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):
        self.dut.raddress_i.value = address
        await FallingEdge(self.dut.clk_i)
        value = self.dut.rdata_o.value
        self.dut.read_i.value = 1
        await FallingEdge(self.dut.clk_i)
        self.dut.read_i.value = 0
        return 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.append(handler)
        self._handle_interrupt = True

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

            if not self._handle_interrupt:
                continue

            if len(self.irq_handlers) == 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 self.irq_handlers:
                    await irq_handler(self)

@cocotb.test()
async def single_transmission(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")

@cocotb.test
async def interrupt(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, 10, "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 = 0,
        div_sel = 0)

    rx = random.randint(0, 255)
    tx = random.randint(0, 255)

    # No interrupt as it's masked
    assert int(dut.interrupt_o.value) == 0
    # TX buffer is empty
    assert int((await driver.read(ADDR_STATUS))) == INTMASK_TX_BUFFER_EMPTY
    await driver.write(ADDR_INTMASK, INTMASK_RX_BUFFER_FULL | INTMASK_TX_BUFFER_EMPTY)
    assert int(dut.interrupt_o.value) == 1

    await slave.send_data(tx, 8)
    await slave.send_data(tx, 8)
    await slave.expect_transaction_in(50, "ns")

    await driver.write(ADDR_DATA, rx)

    async def check_status(value):
        assert int((await driver.read(ADDR_STATUS))) == value
    await cocotb.start(check_status(0))

    await FallingEdge(dut.clk_i)

    assert int(dut.interrupt_o.value) == 0

    timeout = Timer(20, "ns")
    res = await First(RisingEdge(dut.interrupt_o), timeout)

    if res == timeout:
        raise Exception("Timeout on interrupt rising")

    # Check that the buffer is empty, as one transmission should be happening
    assert int((await driver.read(ADDR_STATUS))) == INTMASK_TX_BUFFER_EMPTY
    await driver.write(ADDR_INTMASK, 0)
    assert int(dut.interrupt_o.value) == 0
    assert int((await driver.read(ADDR_STATUS))) == STATUS_TX_BUFFER_EMPTY | STATUS_BUSY

    # Send second data byte
    await driver.write(ADDR_DATA, rx)
    assert int(dut.interrupt_o.value) == 0
    assert int((await driver.read(ADDR_STATUS))) == STATUS_BUSY

    # Now unmask only rx buffer full
    await driver.write(ADDR_INTMASK, INTMASK_RX_BUFFER_FULL)

    timeout = Timer(50, "ns")
    res = await First(timeout, RisingEdge(dut.interrupt_o))

    if res == timeout:
        raise Exception("Timeout on interrupt rising")

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

    assert int(dut.interrupt_o.value) == 1

    await FallingEdge(dut.clk_i)

    assert int(dut.interrupt_o.value) == 0
    assert int((await driver.read(ADDR_STATUS))) == STATUS_TX_BUFFER_EMPTY | STATUS_BUSY

    await slave.wait_all()
    assert int(dut.interrupt_o.value) == 1
    assert int((await driver.read(ADDR_STATUS))) == STATUS_TX_BUFFER_EMPTY | STATUS_RX_BUFFER_FULL

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

@cocotb.test
async def application(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, 10, "ns", csn_pulse = True)
    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 = 0,
        clock_phase = 0,
        pulse_csn = 1,
        lsbfirst = 0,
        size_sel = 1,
        div_sel = 0)

    driver.pending_transactions = 0
    dut_received = []
    async def interrupt_handler(driver):
        await FallingEdge(driver.dut.clk_i)
        status = await driver.read(ADDR_STATUS)
        if int(status) & STATUS_RX_BUFFER_FULL != 0:
            dut_received.append(await driver.read(ADDR_DATA))
            dut._log.info("Received data")

        if int(status) & STATUS_TX_BUFFER_EMPTY != 0 and driver.pending_transactions > 0:
            driver.pending_transactions -= 1
            await driver.write(ADDR_DATA, random.randint(0, 65535))
            dut._log.info("Sent another transaction")

        if int(status) & STATUS_TX_BUFFER_EMPTY != 0 and driver.pending_transactions == 0:
            await driver.write(ADDR_INTMASK, INTMASK_RX_BUFFER_FULL)
            dut._log.info("Turning off tx buffer empty intmask.")

    driver.register_interrupt_handler(interrupt_handler)

    transmitted = []

    # Two transactions.
    driver.pending_transactions = 2
    for i in range(0, driver.pending_transactions):
        data = random.randint(0, 65535)
        transmitted.append(data)
        await slave.send_data(data, 16)
    await slave.expect_transaction_in(50, "ns")
    await driver.write(ADDR_INTMASK, INTMASK_RX_BUFFER_FULL | INTMASK_TX_BUFFER_EMPTY)

    await slave.wait_all()
    assert driver.pending_transactions == 0

    # Wait
    await Timer(100, "ns")

    # Five transactions
    driver.pending_transactions = 5
    for i in range(0, driver.pending_transactions):
        data = random.randint(0, 65535)
        transmitted.append(data)
        await slave.send_data(data, 16)
    await slave.expect_transaction_in(50, "ns")
    await driver.write(ADDR_INTMASK, INTMASK_RX_BUFFER_FULL | INTMASK_TX_BUFFER_EMPTY)

    await slave.wait_all()
    assert driver.pending_transactions == 0
    # Wait
    await Timer(100, "ns")
    # Ten transactions
    driver.pending_transactions = 10
    for i in range(0, driver.pending_transactions):
        data = random.randint(0, 65535)
        transmitted.append(data)
        await slave.send_data(data, 16)
    await slave.expect_transaction_in(100, "ns")
    await driver.write(ADDR_INTMASK, INTMASK_RX_BUFFER_FULL | INTMASK_TX_BUFFER_EMPTY)

    await slave.wait_all()
    assert driver.pending_transactions == 0

    # Everything handled with interrupts


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

    test_filter = os.getenv("TESTCASE", None)
    gui = True if os.getenv("GUI", "0") == "1" else False

    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(
        test_filter = test_filter,
        hdl_toplevel="spi_peripheral", hdl_toplevel_lang=hdl_toplevel_lang,
        test_module="test_spi_peripheral", test_args = extra_args, gui = gui
    )


if __name__ == "__main__":
    spi_peripheral_tests_runner()