from i2c_interfaces import I2CInterface


class TW88xx:

    class Reg:
        SPI_FLASH_MODE_CONTROL_0 = 0X4C0
        SPI_FLASH_MODE_CONTROL_1 = 0X4C1
        DMA_CONTROL = 0X4C3
        FLASH_BUSY_CONTROL = 0X4C4
        WAIT_CONTROL = 0X4C5
        DMA_PAGE = 0X4C6
        DMA_INDEX = 0X4C7
        DMA_LENGTH_MID_BYTE = 0X4C8
        DMA_LENGTH_LOW_BYTE = 0X4C9
        DMA_COMMAND_BUFFER = 0X4CA  # LEN 5
        CLOCK_SWITCH_WAIT_COUNTER_VALUE = 0X4CF
        DMA_READWRITE_BUFFER = 0X4D0  # LEN 8
        SPI_FLASH_STATUS_COMMAND = 0X4D8
        SPI_FLASH_BUSY_CONTROL = 0X4D9
        DMA_LENGTH_HIGH_BYTE = 0X4DA

        LUT_ACCESS_CTRL = 0x410
        LUT_ADDR = 0x411
        LUT_DATA = 0x412

    def __init__(self, i2c: I2CInterface) -> None:
        self.i2c = i2c
        self.page = None

    def _set_page(self, page: int) -> None:
        self.i2c.write(self.ADDR, [0xFF, page])

    def write_reg(self, reg: int, data: int) -> None:
        page = reg >> 8
        if not self.page or self.page != page:
            self._set_page(page)
            self.page = page
        self.i2c.write(self.ADDR, [reg & 0xFF, data])
        # print(f'{reg:03x} = {data:02x}')

    def read_reg(self, reg: int):
        page = reg >> 8
        if not self.page or self.page != page:
            self._set_page(page)
            self.page = page
        self.i2c.write(self.ADDR, [reg & 0xFF])
        return self.i2c.read(self.ADDR, 1)[0]

    def execute_flash_cmd(self, cmd: bytes, read_len=0, read_mode='slow', data=None) -> bytes:
    # def execute_flash_cmd(self, cmd: bytes, read_len=0, read_mode='slow') -> bytes:
        self.write_reg(TW88xx.Reg.SPI_FLASH_MODE_CONTROL_0, 0x00)
        self.write_reg(TW88xx.Reg.DMA_CONTROL, 0x40 + len(cmd))

        self.write_reg(TW88xx.Reg.DMA_PAGE, 0x04)
        self.write_reg(TW88xx.Reg.DMA_INDEX, 0xD0)

        self.write_reg(TW88xx.Reg.DMA_LENGTH_MID_BYTE, 0)
        self.write_reg(TW88xx.Reg.DMA_LENGTH_LOW_BYTE, read_len)
        self.write_reg(TW88xx.Reg.DMA_LENGTH_HIGH_BYTE, 0)

        for i in range(len(cmd)):
            self.write_reg(TW88xx.Reg.DMA_COMMAND_BUFFER + i, cmd[i])

        if data:
            for i in range(len(data)):
                self.write_reg(TW88xx.Reg.DMA_READWRITE_BUFFER + i, data[i])

        self.write_reg(TW88xx.Reg.FLASH_BUSY_CONTROL, 0x01)
        while True:
            busy = self.read_reg(TW88xx.Reg.FLASH_BUSY_CONTROL)
            # print(f'busy: {busy:02x}')
            if (busy & 0x01) == 0:
                break

        data = []
        for i in range(read_len):
            data.append(self.read_reg(TW88xx.Reg.DMA_READWRITE_BUFFER + i))
        return bytes(data)

    def prepare_flash_read(self, cmd_len, read_len):
        self.write_reg(TW88xx.Reg.DMA_CONTROL, 0x40 + cmd_len)

        self.write_reg(TW88xx.Reg.DMA_PAGE, 0x04)
        self.write_reg(TW88xx.Reg.DMA_INDEX, 0xD0)

        self.write_reg(TW88xx.Reg.DMA_LENGTH_MID_BYTE, 0)
        self.write_reg(TW88xx.Reg.DMA_LENGTH_LOW_BYTE, read_len)
        self.write_reg(TW88xx.Reg.DMA_LENGTH_HIGH_BYTE, 0)


    def do_flash_cmd(self, cmd, read_len):
        for i in range(len(cmd)):
            self.write_reg(TW88xx.Reg.DMA_COMMAND_BUFFER + i, cmd[i])

        self.write_reg(TW88xx.Reg.FLASH_BUSY_CONTROL, 0x01)
        while True:
            busy = self.read_reg(TW88xx.Reg.FLASH_BUSY_CONTROL)
            if (busy & 0x01) == 0:
                break

        data = []
        for i in range(read_len):
            data.append(self.read_reg(TW88xx.Reg.DMA_READWRITE_BUFFER + i))
        return bytes(data)

    def dump_regs(self):
        for r in self.REGS:
            val = self.read_reg(r)
            print(f'{r:03x} = {val:02x} ({self.REGS[r]})')

    def dump_lut(self):
        lut_ctrl = self.read_reg(TW88xx.Reg.LUT_ACCESS_CTRL) & 0xF0
        tbl = {}
        for i in range(512):
            data=[]
            for b in range(4):
                self.write_reg(TW88xx.Reg.LUT_ACCESS_CTRL, lut_ctrl | (8 if i>256 else 0) | b)
                self.write_reg(TW88xx.Reg.LUT_ADDR, i & 0xFF)
                _ = self.read_reg(TW88xx.Reg.LUT_DATA)
                data.append(self.read_reg(TW88xx.Reg.LUT_DATA))
            print (f'{i:03x} = {bytes(data).hex()}')
            tbl[i] = data
        return tbl