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

    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, read_len):
        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])
    
    def do_flash_cmd(self, read_len):
        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)