import micropython
from machine import Pin, I2C
from ssd1306 import SSD1306_I2C
import framebuf, sys
import utime

DISP_SDA_PIN = 26
DISP_SCL_PIN = 27
WIDTH = 128
HEIGHT = 64

FPS = 30
TARGET_FRAME_INTERVAL = 1000 // FPS


BUTTON_L_PIN = 17
BUTTON_R_PIN = 16

BUTTON_L = Pin(BUTTON_L_PIN, Pin.IN, Pin.PULL_UP)
BUTTON_R = Pin(BUTTON_R_PIN, Pin.IN, Pin.PULL_UP)
buttons = {
    "l": None,
    "r": None,
}

def irq_button_changed(pin):
    pressed = not pin()
    target = 'l' if pin == BUTTON_L else 'r'
    buttons[target] = utime.ticks_ms() if pressed else None 

    micropython.schedule(lambda _: sm.scene.on_button_change(target, pressed), None)

BUTTON_L.irq(trigger=Pin.IRQ_RISING | Pin.IRQ_FALLING, handler=irq_button_changed)
BUTTON_R.irq(trigger=Pin.IRQ_RISING | Pin.IRQ_FALLING, handler=irq_button_changed)

# def button_r_changed(pin):
#     pressed = not pin()
#     buttons['r'] = utime.ticks_ms() if pressed else None 
# BUTTON_R.irq(trigger=Pin.IRQ_RISING | Pin.IRQ_FALLING, handler=button_r_changed)

class Scene:
    name = 'base_scene'
    def __init__(self, sm):
        self.sm = sm
        self.ctx = sm.ctx    
    
    def on_button_change(self, button, pressed):
        ...

    def init(self):
        ...
    
    def draw(self):
        ...


class MainMenu(Scene):
    name = 'main_menu'

    def on_button_change(self, button, pressed):
        ...
    
    def init(self):
        self.ctx.fill(0)
        self.ctx.rect(0, 0, WIDTH, HEIGHT, 1)
        self.ctx.text('Cool Thing', 12, 12)
        self.ctx.text('Press L and R', 12, 44)
        self.ctx.show()

    def draw(self):
        if buttons['l'] and buttons['r']:
            self.sm.swap(LevelSelect)
            return

        self.ctx.rect(1, 1, 4, HEIGHT-2, 1 if buttons['l'] else 0, True)
        self.ctx.rect(WIDTH - 5, 1, 4, HEIGHT-2, 1 if buttons['r'] else 0, True)
        self.ctx.show()

class LevelSelect(Scene):
    name = 'level_select'

    def init(self):
        self.ctx.fill(0)
        self.ctx.text('hey', 0, 0)
        self.ctx.show()

def init_i2c(scl_pin, sda_pin):
    """find the i2c oled"""
    i2c_dev = I2C(1, scl=Pin(scl_pin), sda=Pin(sda_pin), freq=1600000)
    i2c_addr = [hex(ii) for ii in i2c_dev.scan()]
    
    if not i2c_addr:
        print('No I2C Display Found')
        sys.exit()
    else:
        print("I2C Address      : {}".format(i2c_addr[0]))
        print("I2C Configuration: {}".format(i2c_dev))
    
    return i2c_dev

class StateMachine:
    def __init__(self):
        self.ctx = SSD1306_I2C(
            WIDTH, HEIGHT, init_i2c(scl_pin=DISP_SCL_PIN, 
                                    sda_pin=DISP_SDA_PIN))

        self.last_frame_ms = utime.ticks_ms()
        self.last_log_ms = utime.ticks_ms()
        self.last_log_frames = 0

        self.scene = MainMenu(self)
        self.scene.init()

    def swap(self, NewScene, *args, **kwargs):
        self.scene = NewScene(self, *args, **kwargs)
        self.scene.init()
    
    def run_forever(self):
        while True:
            self.scene.draw()
            self.wait_for_frame()
    
    def wait_for_frame(self):
        self.last_log_frames += 1
        now = utime.ticks_ms()
        if utime.ticks_diff(now, self.last_log_ms) > 1000:
            print(f'FPS: {self.last_log_frames} / {FPS}')
            self.last_log_ms = now
            self.last_log_frames = 0

        spent_ms = utime.ticks_diff(now, self.last_frame_ms)
        remaining = int(TARGET_FRAME_INTERVAL - spent_ms)

        if remaining > 0:
            utime.sleep_ms(remaining)
            self.last_frame_ms = utime.ticks_ms()
        else:
            self.last_frame_ms = now

sm = StateMachine()
if __name__ == '__main__':
    sm.run_forever()