from machine import Pin, SPI
import framebuf
import time
import BasicFont

# Pin Definitions
DC = 8
RST = 12
MOSI = 11
SCK = 10
CS = 9

KEY0 = 15
KEY1 = 17

# Display object
display = None


class OLED_1inch3_SPI(framebuf.FrameBuffer):
    def __init__(self):
        self.is_on = 0
        self.width = 128
        self.height = 64

        self.white = 0xffff
        self.black = 0x0000

        self.font = BasicFont.BasicFontCondensed

        # framebuf init
        self.buffer = bytearray(self.height * self.width // 8)
        super().__init__(self.buffer, self.width, self.height, framebuf.MONO_HMSB)

        # SPI init
        self.cs = Pin(CS, Pin.OUT)
        self.rst = Pin(RST, Pin.OUT)

        self.cs(1)
        self.spi = SPI(1, 20000_000, polarity=0, phase=0, sck=Pin(SCK), mosi=Pin(MOSI), miso=None)
        self.dc = Pin(DC, Pin.OUT)
        self.dc(1)

        # display init
        self.init_display()

        # Clear the screen on init - needs framebuf init, spi init and display init
        self.clear()

        # Init Keys
        self.KEY0 = KEY0
        self.KEY1 = KEY1
        self.key0 = Pin(KEY0, Pin.IN, Pin.PULL_UP)
        self.key1 = Pin(KEY1, Pin.IN, Pin.PULL_UP)

    def is_pressed(self, key):
        if key == self.KEY0:
            return not self.key0.value()
        elif key == self.KEY1:
            return not self.key1.value()
        else:
            return None

    def on(self):
        if not self.is_on:
            self.write_cmd(0xAF)
            self.is_on = 1

    def off(self):
        if self.is_on:
            self.write_cmd(0xAE)
            self.is_on = 0

    def get_width(self):
        return self.width

    def get_height(self):
        return self.height

    def write_cmd(self, cmd):
        self.cs(1)
        self.dc(0)
        self.cs(0)
        self.spi.write(bytearray([cmd]))
        self.cs(1)

    def write_data(self, buf):
        self.cs(1)
        self.dc(1)
        self.cs(0)
        self.spi.write(bytearray([buf]))
        self.cs(1)

    def init_display(self):
        """Initialize display"""
        self.rst(1)
        time.sleep(0.001)
        self.rst(0)
        time.sleep(0.01)
        self.rst(1)

        self.off()

        self.write_cmd(0x00)  # set lower column address
        self.write_cmd(0x10)  # set higher column address

        self.write_cmd(0xB0)  # set page address

        self.write_cmd(0xdc)  # set display start line
        self.write_cmd(0x00)  # (2nd param)

        self.write_cmd(0x81)  # contract control
        self.write_cmd(0x80)  # 128

        self.write_cmd(0x21)  # Set Memory addressing mode (0x20/0x21)
        self.write_cmd(0xa0)  # set segment remap
        self.write_cmd(0xc0)  # Com scan direction
        self.write_cmd(0xa4)  # Disable Entire Display On (0xA4/0xA5)
        self.write_cmd(0xa6)  # normal / reverse

        self.write_cmd(0xa8)  # multiplex ratio ??
        self.write_cmd(0x3f)  # duty = 1/64 ??

        self.write_cmd(0xd3)  # set display offset
        self.write_cmd(0x60)

        self.write_cmd(0xd5)  # set osc division
        self.write_cmd(0x50)

        self.write_cmd(0xd9)  # set pre-charge period
        self.write_cmd(0x22)

        self.write_cmd(0xdb)  # set vcomh
        self.write_cmd(0x35)

        self.write_cmd(0xad)  # set charge pump enable
        self.write_cmd(0x8a)  # Set DC-DC enable (a=0:disable; a=1:enable)

        self.on()

    # --- Helper functions for scaling bitmaps ---
    def _unpack_bitmap(self, data, width, height):
        """
        Unpack a MONO_VLSB bitmap into a 2D list of 0s and 1s.
        Each byte in data represents one vertical column of 8 pixels.
        """
        arr = [[0] * width for _ in range(height)]
        for x in range(width):
            b = data[x]
            for y in range(height):
                arr[y][x] = 1 if (b >> y) & 1 else 0
        return arr

    def _scale_2d_array(self, arr, scale):
        """
        Scale a 2D list (of 0s and 1s) by an integer factor using pixel replication.
        """
        h = len(arr)
        w = len(arr[0])
        new_arr = [[0] * (w * scale) for _ in range(h * scale)]
        for y in range(h):
            for x in range(w):
                pixel = arr[y][x]
                for dy in range(scale):
                    for dx in range(scale):
                        new_arr[y * scale + dy][x * scale + dx] = pixel
        return new_arr

    def _pack_bitmap(self, arr):
        """
        Pack a 2D list of 0s and 1s into a MONO_VLSB bytearray.
        The height must be a multiple of 8.
        """
        height = len(arr)
        width = len(arr[0])
        buf = bytearray(width * (height // 8))
        for x in range(width):
            for byte_row in range(height // 8):
                b = 0
                for bit in range(8):
                    y = byte_row * 8 + bit
                    if arr[y][x]:
                        b |= (1 << bit)
                buf[x + byte_row * width] = b
        return buf

    def _scale_bitmap(self, data, width, height, scale):
        """
        Scale a MONO_VLSB bitmap (provided as a bytearray) by an integer factor.
        Returns the scaled bytearray and the new dimensions.
        """
        arr = self._unpack_bitmap(data, width, height)
        scaled_arr = self._scale_2d_array(arr, scale)
        new_buffer = self._pack_bitmap(scaled_arr)
        new_width = width * scale
        new_height = height * scale
        return new_buffer, new_width, new_height

    # --- Modified text() method with scaling support ---
    # The original text is assumed to be 8 pixels high.
    def text(self, s, x0, y0, col=0xffff, wrap=1, just=0, scale=1):
        """
        Draw text on the display starting at (x0, y0).

        Parameters:
          s      : The string to render.
          x0, y0 : Starting coordinates.
          col    : Color (default white). If 0, the text is inverted.
          wrap   : Text wrapping mode.
                   0: Clip at the right edge.
                   1: Wrap to the next line.
                   Other: Wrap to the original x0 coordinate.
          just   : Justification.
                   0: Left-justified.
                   1: Right-justified.
                   2: Center-justified.
          scale  : Scaling factor for enlarging the font (default is 1, i.e. no scaling).

        Returns:
          A list [new_x, new_y] with updated coordinates after rendering the text.
        """
        if len(s) == 0:
            return (x0, y0)

        x = x0
        pixels = bytearray([])
        # For each character in the string, accumulate its bitmap data
        for i in range(len(s)):
            C = ord(s[i])
            if C < 32 or C > 127:
                C = 32
            cdata = self.font[C - 32]
            effective_char_width = len(cdata) * scale
            effective_pixels_width = len(pixels) * scale
            # Check if adding the next character would exceed the display width
            if len(pixels) and (
                    (just == 0 and x + effective_pixels_width + effective_char_width > self.width) or
                    (just == 1 and x - effective_pixels_width - effective_char_width < 0) or
                    (just == 2 and (x - effective_pixels_width / 2 - effective_char_width < 0 or
                                    x + effective_pixels_width / 2 + effective_char_width > self.width))
            ):
                # Invert pixels if needed
                if col == 0:
                    for j, v in enumerate(pixels):
                        pixels[j] = 0xFF & ~v
                # Create a frame buffer from the accumulated data; scale if needed
                if scale == 1:
                    fb = framebuf.FrameBuffer(pixels, len(pixels), 8, framebuf.MONO_VLSB)
                else:
                    scaled_buf, new_w, new_h = self._scale_bitmap(pixels, len(pixels), 8, scale)
                    fb = framebuf.FrameBuffer(scaled_buf, new_w, new_h, framebuf.MONO_VLSB)
                if just == 0:
                    self.blit(fb, x, y0)
                elif just == 1:
                    self.blit(fb, x - (len(pixels) * scale), y0)
                else:
                    self.blit(fb, x - int((len(pixels) * scale) / 2), y0)
                pixels = bytearray([])
                if wrap == 0:
                    return [x, y0 + (8 * scale) + 1]
                if wrap == 1:
                    x = 0
                else:
                    x = x0
                y0 = y0 + (8 * scale) + 1
                if y0 > self.height:
                    return [x, y0]
            # Accumulate character bitmap data (unscaled)
            pixels += bytearray(cdata)

        # Render any remaining accumulated text
        if col == 0:
            for j, v in enumerate(pixels):
                pixels[j] = 0xFF & ~v
        if scale == 1:
            fb = framebuf.FrameBuffer(pixels, len(pixels), 8, framebuf.MONO_VLSB)
        else:
            scaled_buf, new_w, new_h = self._scale_bitmap(pixels, len(pixels), 8, scale)
            fb = framebuf.FrameBuffer(scaled_buf, new_w, new_h, framebuf.MONO_VLSB)
        if just == 0:
            self.blit(fb, x, y0)
        elif just == 1:
            self.blit(fb, x - (len(pixels) * scale), y0)
        else:
            self.blit(fb, x - int((len(pixels) * scale) / 2), y0)
        return [x, y0 + (8 * scale) + 1]

    # Shows the framebuffer contents on the display.
    # If no arguments are given, the full frame buffer is sent to the display.
    # startXPage:  The horizontal page index to start the update. The X-axis is divided in 16 8-pixel 'pages' (0 ~ 15)
    # startYLine:  The vertical line index to start the display update (0 ~ 63)
    # endXPage:    The horizontal page index to end the update (excluding that page index) (1 ~ 16)
    # endYLine:    The vertical line index to end the display update (excluding that line index) (1 ~ 64)
    def show(self, startXPage=0, startYLine=0, endXPage=16, endYLine=64):
        self.__validateShowArguments(startXPage, startYLine, endXPage, endYLine)
        doCustomPageAddressing = startXPage > 0 or endXPage < 16
        if not doCustomPageAddressing:
            self.write_cmd(0xB0)
        for yLine in range(startYLine, endYLine):
            columnSramAddress = 63 - yLine
            self.write_cmd(0x00 + (columnSramAddress & 0x0f))
            self.write_cmd(0x10 + (columnSramAddress >> 4))
            if doCustomPageAddressing:
                self.write_cmd(0xB0 + (startXPage & 0x0f))
            for num in range(startXPage, endXPage):
                self.write_data(self.buffer[yLine * 16 + num])

    def __validateShowArguments(self, startXPage, startYLine, endXPage, endYLine):
        if not startYLine < endYLine:
            raise IndexError("show(...): startYLine (" + str(startYLine) +
                             ") must be smaller than endYLine (" + str(endYLine) + ").")
        if startYLine < 0 or startYLine > 63:
            raise IndexError("show(...): startYLine acceptable range is 0 ~ 63. Given: " + str(startYLine))
        if endYLine < 1 or endYLine > 64:
            raise IndexError("show(...): endYLine acceptable range is 1 ~ 64. Given: " + str(endYLine))
        if not startXPage < endXPage:
            raise IndexError("show(...): startXPage (" + str(startXPage) +
                             ") must be smaller than endXPage (" + str(endXPage) + ").")
        if startXPage < 0 or startXPage > 15:
            raise IndexError("show(...): startXPage acceptable range is 0 ~ 15. Given: " + str(startXPage))
        if endXPage < 1 or endXPage > 16:
            raise IndexError("show(...): endXPage acceptable range is 1 ~ 16. Given: " + str(endXPage))

    def auto_text(self, s, col=0xffff):
        """
        Automatically determine the largest integer scale factor such that the full
        string (with whole words kept on the same line) fits on the screen,
        and render it centered.

        Parameters:
          s   : The text string to render.
          col : Color (default white).
        """
        # Maximum possible scale based on unscaled height (8 pixels per line)
        max_possible_scale = self.height // 8
        space_width = 3  # unscaled space width in pixels

        # Try scales from the largest down to 1
        for scale in range(max_possible_scale, 0, -1):
            words = s.split(' ')
            scale_valid = True
            # Check that every individual word fits on one line at this scale
            for word in words:
                # Compute the word width using the font data
                word_width = sum(len(self.font[ord(c) - 32]) * scale for c in word if 32 <= ord(c) <= 127)
                if word_width > self.width:
                    scale_valid = False
                    break
            if not scale_valid:
                continue  # Try a smaller scale

            # Build lines without breaking words.
            lines = []
            current_line_words = []
            current_line_width = 0
            for word in words:
                # Calculate word width at current scale.
                word_width = sum(len(self.font[ord(c) - 32]) * scale for c in word if 32 <= ord(c) <= 127)
                additional_space = space_width * scale if current_line_words else 0
                if current_line_words and (current_line_width + additional_space + word_width > self.width):
                    # Commit current line and start a new one
                    lines.append(" ".join(current_line_words))
                    current_line_words = [word]
                    current_line_width = word_width
                else:
                    if current_line_words:
                        current_line_width += additional_space + word_width
                    else:
                        current_line_width += word_width
                    current_line_words.append(word)
            if current_line_words:
                lines.append(" ".join(current_line_words))

            # Check total height: each line is (8*scale) pixels high with a 1-pixel gap between lines.
            total_height = len(lines) * (8 * scale) + (len(lines) - 1)
            if total_height <= self.height:
                # This scale works
                break

        # Center the text vertically.
        total_text_height = len(lines) * (8 * scale) + (len(lines) - 1)
        y_offset = (self.height - total_text_height) // 2

        self.clear()

        # Render each line centered horizontally.
        for line in lines:
            line_width = 0
            for char in line:
                if char == ' ':
                    line_width += space_width * scale
                else:
                    C = ord(char)
                    if C < 32 or C > 127:
                        C = 32
                    line_width += len(self.font[C - 32]) * scale
            x_offset = (self.width - line_width) // 2
            self.text(line, x_offset, y_offset, col, wrap=0, just=0, scale=scale)
            y_offset += (8 * scale) + 1

    def clear(self):
        self.fill(self.black)
        self.show()



def get():
    global display
    if display is None:
        display = OLED_1inch3_SPI()
    return display


def test():
    display = get()
    display.clear()
    print("Running display tests")

    # Test sequence to validate Xpage and YLine writing to the display
    display.fill(1)
    display.show()
    display.fill(0)
    display.show()

    # Fill the framebuffer with white in specific regions:
    display.fill(1)
    display.show(0, 0, 16, 8)  # Top horizontal bar
    display.show(0, 56, 16, 64)  # Bottom horizontal bar
    display.show(0, 8, 1, 56)  # Vertical bar on the left, partially drawn
    display.show(7, 24, 9, 40)  # 16 x 16 pixel square in the middle


if __name__ == '__main__':
    test()
