"""
Waveshare 2.9" e-Paper V2 driver for MicroPython on Raspberry Pi Pico.

V2 adaptation of the original V1 driver.
Compatible with Waveshare epd2in9_V2.py command flow.

Controller:
    SSD1680 / UC8151-compatible command set used by Waveshare V2

This version intentionally keeps the original V1 architecture
and only adapts:
    - init sequence
    - LUT handling
    - refresh/update flow
    - busy handling

API compatibility:
    EPD(spi, cs, dc, rst, busy)
    init()
    clear()
    draw_text(x, y, text)
    show()
    sleep()
"""

from micropython import const
from time import sleep_ms, ticks_ms, ticks_diff
import ustruct
import framebuf

EPD_WIDTH = const(128)
EPD_HEIGHT = const(296)

# V2 modules are usually busy=1 while processing
BUSY = const(1)

# Commands (same family as original V1 driver)
DRIVER_OUTPUT_CONTROL = const(0x01)
BOOSTER_SOFT_START_CONTROL = const(0x0C)
DEEP_SLEEP_MODE = const(0x10)
DATA_ENTRY_MODE_SETTING = const(0x11)
SW_RESET = const(0x12)

MASTER_ACTIVATION = const(0x20)
DISPLAY_UPDATE_CONTROL_1 = const(0x21)
DISPLAY_UPDATE_CONTROL_2 = const(0x22)

WRITE_RAM = const(0x24)
WRITE_VCOM_REGISTER = const(0x2C)
WRITE_LUT_REGISTER = const(0x32)

SET_DUMMY_LINE_PERIOD = const(0x3A)
SET_GATE_TIME = const(0x3B)

BORDER_WAVEFORM_CONTROL = const(0x3C)

SET_RAM_X_ADDRESS_START_END_POSITION = const(0x44)
SET_RAM_Y_ADDRESS_START_END_POSITION = const(0x45)
SET_RAM_X_ADDRESS_COUNTER = const(0x4E)
SET_RAM_Y_ADDRESS_COUNTER = const(0x4F)

TERMINATE_FRAME_READ_WRITE = const(0xFF)


class EPDTimeout(Exception):
    pass


class EPD:

    # LUT from Waveshare epd2in9_V2.py
    LUT_FULL_UPDATE = bytearray([
        0x80,0x60,0x40,0x00,0x00,0x10,0x60,0x20,
        0x00,0x00,0x80,0x60,0x40,0x00,0x00,0x10,
        0x60,0x20,0x00,0x00,0x00,0x00,0x00,0x00,
        0x00,0x00,0x00,0x00,0x00,0x00
    ])

    LUT_PARTIAL_UPDATE = bytearray([
        0x00,0x40,0x00,0x00,0x00,0x10,0x00,0x00,
        0x00,0x00,0x80,0x80,0x00,0x00,0x00,0x20,
        0x20,0x00,0x00,0x00,0x00,0x00,0x00,0x00,
        0x00,0x00,0x00,0x00,0x00,0x00
    ])

    def __init__(self, spi, cs, dc, rst, busy):

        self.spi = spi
        self.cs = cs
        self.dc = dc
        self.rst = rst
        self.busy = busy

        self.cs.init(self.cs.OUT, value=1)
        self.dc.init(self.dc.OUT, value=0)
        self.rst.init(self.rst.OUT, value=0)
        self.busy.init(self.busy.IN)

        self.width = EPD_WIDTH
        self.height = EPD_HEIGHT

        self.buffer = bytearray(self.width * self.height // 8)

        self.fb = framebuf.FrameBuffer(
            self.buffer,
            self.width,
            self.height,
            framebuf.MONO_HLSB
        )

        self.clear()

    # ---------------------------------------------------------
    # low level SPI
    # ---------------------------------------------------------

    def _command(self, command, data=None):

        self.dc(0)
        self.cs(0)
        self.spi.write(bytearray([command]))
        self.cs(1)

        if data is not None:
            self._data(data)

    def _data(self, data):

        self.dc(1)
        self.cs(0)

        if isinstance(data, int):
            self.spi.write(bytearray([data]))
        else:
            self.spi.write(data)

        self.cs(1)

    # ---------------------------------------------------------
    # busy handling
    # ---------------------------------------------------------

    def wait_until_idle(self, timeout_ms=15000):

        start = ticks_ms()

        while self.busy.value() == BUSY:
            if ticks_diff(ticks_ms(), start) > timeout_ms:
                raise EPDTimeout("e-Paper BUSY timeout")

            sleep_ms(20)

    # ---------------------------------------------------------
    # reset
    # ---------------------------------------------------------

    def reset(self):

        self.rst(0)
        sleep_ms(200)

        self.rst(1)
        sleep_ms(200)

    # ---------------------------------------------------------
    # LUT
    # ---------------------------------------------------------

    def set_lut(self, lut):

        self._command(WRITE_LUT_REGISTER, lut)

    # ---------------------------------------------------------
    # init
    # ---------------------------------------------------------

    def init(self):

        self.reset()

        self.wait_until_idle()

        self._command(SW_RESET)
        self.wait_until_idle()

        # Driver output control
        self._command(
            DRIVER_OUTPUT_CONTROL,
            ustruct.pack("<HB", EPD_HEIGHT - 1, 0x00)
        )

        # data entry mode
        self._command(DATA_ENTRY_MODE_SETTING, b'\x03')

        # RAM area
        self.set_memory_area(
            0,
            0,
            self.width - 1,
            self.height - 1
        )

        self.set_memory_pointer(0, 0)

        # border waveform
        self._command(BORDER_WAVEFORM_CONTROL, b'\x05')

        # display update
        self._command(DISPLAY_UPDATE_CONTROL_1, b'\x00\x80')

        # temperature / timing
        self._command(SET_DUMMY_LINE_PERIOD, b'\x1A')
        self._command(SET_GATE_TIME, b'\x08')

        # VCOM
        self._command(WRITE_VCOM_REGISTER, b'\x26')

        # booster
        self._command(
            BOOSTER_SOFT_START_CONTROL,
            b'\xD7\xD6\x9D'
        )

        # LUT
        self.set_lut(self.LUT_FULL_UPDATE)

    # ---------------------------------------------------------
    # memory handling
    # ---------------------------------------------------------

    def set_memory_area(self, x_start, y_start, x_end, y_end):

        self._command(
            SET_RAM_X_ADDRESS_START_END_POSITION
        )

        self._data(
            bytearray([
                (x_start >> 3) & 0xFF,
                (x_end >> 3) & 0xFF
            ])
        )

        self._command(
            SET_RAM_Y_ADDRESS_START_END_POSITION,
            ustruct.pack("<HH", y_start, y_end)
        )

    def set_memory_pointer(self, x, y):

        self._command(SET_RAM_X_ADDRESS_COUNTER)
        self._data(bytearray([(x >> 3) & 0xFF]))

        self._command(
            SET_RAM_Y_ADDRESS_COUNTER,
            ustruct.pack("<H", y)
        )

        self.wait_until_idle()

    # ---------------------------------------------------------
    # panel clear
    # ---------------------------------------------------------

    def clear_frame_memory(self, color):

        self.set_memory_area(
            0,
            0,
            self.width - 1,
            self.height - 1
        )

        self.set_memory_pointer(0, 0)

        self._command(WRITE_RAM)

        for _ in range(self.width // 8 * self.height):
            self._data(color)

    # ---------------------------------------------------------
    # refresh
    # ---------------------------------------------------------

    def display_frame(self):

        # V2 refresh sequence
        self._command(DISPLAY_UPDATE_CONTROL_2, b'\xF7')

        self._command(MASTER_ACTIVATION)

        self._command(TERMINATE_FRAME_READ_WRITE)

        self.wait_until_idle()

    # ---------------------------------------------------------
    # public API
    # ---------------------------------------------------------

    def clear(self):

        self.fb.fill(1)

    def fill(self, color):

        self.fb.fill(1 if color else 0)

    def draw_text(self, x, y, text):

        self.fb.text(str(text), int(x), int(y), 0)

    def show(self):

        self.set_memory_area(
            0,
            0,
            self.width - 1,
            self.height - 1
        )

        self.set_memory_pointer(0, 0)

        self._command(WRITE_RAM)

        self._data(self.buffer)

        self.display_frame()

    def hard_clear(self):

        self.clear_frame_memory(0xFF)

        self.display_frame()

    # ---------------------------------------------------------
    # sleep
    # ---------------------------------------------------------

    def sleep(self):

        self._command(DEEP_SLEEP_MODE, b'\x01')

        sleep_ms(100)

    # ---------------------------------------------------------
    # drawing helpers
    # ---------------------------------------------------------

    def pixel(self, x, y, color=0):

        self.fb.pixel(int(x), int(y), color)

    def hline(self, x, y, w, color=0):

        self.fb.hline(int(x), int(y), int(w), color)

    def vline(self, x, y, h, color=0):

        self.fb.vline(int(x), int(y), int(h), color)

    def line(self, x1, y1, x2, y2, color=0):

        self.fb.line(
            int(x1),
            int(y1),
            int(x2),
            int(y2),
            color
        )

    def rect(self, x, y, w, h, color=0):

        self.fb.rect(
            int(x),
            int(y),
            int(w),
            int(h),
            color
        )

    def fill_rect(self, x, y, w, h, color=0):

        self.fb.fill_rect(
            int(x),
            int(y),
            int(w),
            int(h),
            color
        )

    def draw_text_scaled(
        self,
        x,
        y,
        text,
        scale=2,
        color=0
    ):

        text = str(text)

        if not text:
            return

        if scale <= 1:
            self.fb.text(
                text,
                int(x),
                int(y),
                color
            )
            return

        char_count = len(text)

        src_w = char_count * 8

        stride = (src_w + 7) >> 3

        tmp = bytearray(stride * 8)

        tmp_fb = framebuf.FrameBuffer(
            tmp,
            src_w,
            8,
            framebuf.MONO_HLSB
        )

        tmp_fb.fill(1)

        tmp_fb.text(text, 0, 0, 0)

        x0 = int(x)
        y0 = int(y)

        for ty in range(8):
            for tx in range(src_w):

                if tmp_fb.pixel(tx, ty) == 0:

                    self.fb.fill_rect(
                        x0 + tx * scale,
                        y0 + ty * scale,
                        scale,
                        scale,
                        color,
                    )