from machine import Pin, SoftI2C
from ssd1306 import SSD1306_I2C
import time

class Screen:

    def __init__(self):
        # OLED pins on Heltec V3
        self.sda = Pin(17)
        self.scl = Pin(18)
        self.rst = Pin(21, Pin.OUT)

        self.__reset_oled()

        self.i2c = SoftI2C(scl=self.scl, sda=self.sda, freq=100000)
        self.oled = SSD1306_I2C(128, 64, self.i2c, addr=0x3C)

        self.clear()

    def __reset_oled(self):
        self.rst.value(0)
        time.sleep_ms(150)
        self.rst.value(1)
        time.sleep_ms(150)

    def clear(self):
        self.oled.fill(0)
        self.oled.show()

    def display(self, message, col=0, row=0):
        self.clear()
        self.oled.text(message, col, row)
        self.oled.show()

    def display_lines(self, lines, col):
        self.clear()
        for i, line in enumerate(lines):
            self.oled.text(line, col, i * 20)
        self.oled.show()