import struct
import cocotb
from cocotb.clock import Clock
from cocotb.triggers import RisingEdge, FallingEdge, Timer
from cocotb_bus.drivers import BusDriver


class SPIRegMap(BusDriver):
    _signals = ["SPI_MISO", "SPI_MOSI", "SPI_CS", "SPI_SClk"]

    def __init__(self, instr, name=None, freq=1e6, **kwargs):
        BusDriver.__init__(self, instr, name, instr.clk, **kwargs)
        self.freq = round(1e9 / freq, 4)

        # Drive some sensible defaults (setimmediatevalue to avoid x asserts)
        self.bus.SPI_MOSI.setimmediatevalue(0)
        self.bus.SPI_CS.setimmediatevalue(1)
        self.bus.SPI_SClk.setimmediatevalue(0)

    async def writereg(self, reg, value):
        b = struct.pack('<II', (reg << 2) | 0x80000000, value)
        res = await self.spi_xfer2(b)

    async def writeregs(self, regs):
        for r, v, in regs:
            b = struct.pack('<II', (r << 2) | 0x80000000, v)
            await self.spi_xfer2(b)

    async def readreg(self, reg):
        b = struct.pack('<II', reg << 2, 0)
        res = await self.spi_xfer2(b)
        return struct.unpack('<II', bytes(res))[1]

    async def writeblock(self, addr, data):
        b = struct.pack('<I', addr | 0x80000000) + data
        res = await self.spi_xfer2(b)

    async def readblock(self, addr, length):
        b = struct.pack('<I', addr) + bytes([0] * length)
        res = bytes(await self.spi_xfer2(b))
        return res[4:]

    async def spi_xfer2(self, data):
        """ should look like SpiDev.xfer2()
        """
        self.bus.SPI_CS.value = 0
        self.bus.SPI_MOSI.value = (data[0] & 0x80) >> 7

        # this timing changes wildly between spi controllers
        await Timer(time=0.3, units='us', round_mode='round')
        sclk_task = cocotb.start_soon(
            Clock(self.bus.SPI_SClk, self.freq, units='ns').start(cycles=len(data) * 8))

        miso_data = []

        for d in data:
            m = 0
            for i in range(8):
                self.bus.SPI_MOSI.value = (d & 0x80) >> 7
                d = d << 1
                await RisingEdge(self.bus.SPI_SClk)
                m = (m << 1) | self.bus.SPI_MISO.value
                await FallingEdge(self.bus.SPI_SClk)
            miso_data.append(m)

        await sclk_task
        await Timer(time=0.3, units='us', round_mode='round')
        self.bus.SPI_CS.value = 1
        await Timer(time=0.3, units='us', round_mode='round')
        return miso_data
