from machine import I2C
import framebuf

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

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

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

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

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

    def show(self):
        raise NotImplementedError

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

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

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

    def init_display(self):
        cmds = [
            0xAE,  # DISPLAY_OFF
            0xD5, 0x80,  # SET_DISPLAY_CLOCK_DIV
            0xA8, self.height - 1,  # SET_MULTIPLEX
            0xD3, 0x00,  # SET_DISPLAY_OFFSET
            0x40,  # SET_START_LINE
            0x8D, 0x14 if not self.external_vcc else 0x10,  # CHARGE_PUMP
            0x20, 0x00,  # MEMORY_MODE
            0xA1,  # SEG_REMAP
            0xC8,  # COM_SCAN_DEC
            0xDA, 0x12,  # SET_COM_PINS
            0x81, 0xCF,  # SET_CONTRAST
            0xD9, 0xF1 if not self.external_vcc else 0x22,  # SET_PRECHARGE
            0xDB, 0x40,  # SET_VCOM_DETECT
            0xA4,  # DISPLAY_ALL_ON_RESUME
            0xA6,  # NORMAL_DISPLAY
            0xAF  # DISPLAY_ON
        ]
        for cmd in cmds:
            self.write_cmd(cmd)
        self.fill(0)
        self.show()

    def show(self):
        self.write_cmd(0x21)  # COLUMN_ADDR
        self.write_cmd(0)     # start column
        self.write_cmd(self.width - 1)  # end column
        self.write_cmd(0x22)  # PAGE_ADDR
        self.write_cmd(0)     # start page
        self.write_cmd(self.pages - 1)  # end page
        self.write_data(self.buffer)
