import framebuf

_SET_CONTRASTE      = 0x81
_SET_ENTRY_ON       = 0xA4
_SET_NORM_INV       = 0xA6
_SET_DISP           = 0xAE
_SET_MEM_ADDR       = 0x20
_SET_COL_ADDR       = 0x21
_SET_PAGE_ADDR      = 0x22
_SET_DISP_START     = 0x40
_SET_SEG_REMAP      = 0xA0
_SET_MUX_RATIO      = 0xA8
_SET_COM_OUT_DIR    = 0xC0
_SET_DISP_OFFSET    = 0xD3
_SET_COM_PIN_CFG    = 0xDA
_SET_DISP_CLK_DIV   = 0xD5
_SET_PRECHARGE      = 0xD9
_SET_VCOM_DESEL     = 0xDB
_SET_CHARGE_PUMP    = 0x8D


class SSD1306_I2C(framebuf.FrameBuffer):

    def __init__(self, largura, altura, i2c, endereco=0x3C):
        self.largura = largura
        self.altura = altura
        self.i2c = i2c
        self.endereco = endereco
        self.paginas = altura // 8
        self.buffer = bytearray(self.paginas * largura)
        super().__init__(self.buffer, largura, altura, framebuf.MONO_VLSB)
        self._comando_buf = bytearray(2)
        self.iniciar()

    def comando(self, valor):
        self._comando_buf[0] = 0x80
        self._comando_buf[1] = valor
        self.i2c.writeto(self.endereco, self._comando_buf)

    def iniciar(self):
        sequencia = (
            _SET_DISP,
            _SET_MEM_ADDR, 0x00,
            _SET_DISP_START,
            _SET_SEG_REMAP | 0x01,
            _SET_MUX_RATIO, self.altura - 1,
            _SET_COM_OUT_DIR | 0x08,
            _SET_DISP_OFFSET, 0x00,
            _SET_COM_PIN_CFG, 0x02 if self.altura == 32 else 0x12,
            _SET_DISP_CLK_DIV, 0x80,
            _SET_PRECHARGE, 0xF1,
            _SET_VCOM_DESEL, 0x30,
            _SET_CONTRASTE, 0xFF,
            _SET_ENTRY_ON,
            _SET_NORM_INV,
            _SET_CHARGE_PUMP, 0x14,
            _SET_DISP | 0x01,
        )
        for c in sequencia:
            self.comando(c)
        self.fill(0)
        self.show()

    def contraste(self, valor):
        self.comando(_SET_CONTRASTE)
        self.comando(valor)

    def show(self):
        self.comando(_SET_COL_ADDR)
        self.comando(0)
        self.comando(self.largura - 1)
        self.comando(_SET_PAGE_ADDR)
        self.comando(0)
        self.comando(self.paginas - 1)
        self.i2c.writeto(self.endereco, b"\x40" + self.buffer)