# MicroPython SSD1306 OLED driver, I2C and SPI interfaces

import framebuf

class SSD1306:
    def __init__(self, width, height):
        self.width = width
        self.height = height
        self.pages = height // 8
        self.buffer = bytearray(self.pages * width)
        self.framebuf = framebuf.FrameBuffer(self.buffer, width, height, framebuf.MONO_VLSB)
        self.poweron()
        self.init_display()

    def init_display(self):
        for cmd in (
            0xae,         # DISPLAYOFF
            0x20, 0x00,   # MEMORYMODE
            0x40,         # SETSTARTLINE
            0xa1,         # SEGREMAP
            0xc8,         # COMSCANDEC
            0xda, 0x12,   # SETCOMPINS
            0x81, 0x7f,   # SETCONTRAST
            0xa4,         # DISPLAYALLON_RESUME
            0xa6,         # NORMALDISPLAY
            0xd5, 0x80,   # SETDISPLAYCLOCKDIV
            0x8d, 0x14,   # CHARGEPUMP
            0xaf          # DISPLAYON
        ):
            self.write_cmd(cmd)

    def poweron(self):
        pass

    def contrast(self, contrast):
        self.write_cmd(0x81)
        self.write_cmd(contrast)

    def invert(self, invert):
        self.write_cmd(0xa7 if invert else 0xa6)

    def show(self):
        for page in range(self.pages):
            self.write_cmd(0xb0 | page)
            self.write_cmd(0x00)
            self.write_cmd(0x10)
            self.write_data(self.buffer[page * self.width:(page + 1) * self.width])

    def fill(self, col):
        self.framebuf.fill(col)

    def pixel(self, x, y, col):
        self.framebuf.pixel(x, y, col)

    def text(self, string, x, y):
        self.framebuf.text(string, x, y)

    def scroll(self, dx, dy):
        self.framebuf.scroll(dx, dy)

    def blit(self, fbuf, x, y):
        self.framebuf.blit(fbuf, x, y)


class SSD1306_I2C(SSD1306):
    def __init__(self, width, height, i2c, addr=0x3c):
        self.i2c = i2c
        self.addr = addr
        super().__init__(width, height)

    def write_cmd(self, cmd):
        self.i2c.writeto(self.addr, bytearray([0x00, cmd]))

    def write_data(self, buf):
        self.i2c.writeto(self.addr, bytearray([0x40]) + buf)