from machine import Pin, I2C, PWM
import time
import random
from ssd1306 import SSD1306_I2C

WIDTH = 128
HEIGHT = 64

# Buttons
button1 = Pin(7, Pin.IN, Pin.PULL_UP)
button2 = Pin(8, Pin.IN, Pin.PULL_UP)
button3 = Pin(9, Pin.IN, Pin.PULL_UP)

# LEDs with PWM for dimming
led1 = PWM(Pin(20))
led2 = PWM(Pin(21))
led3 = PWM(Pin(22))

led1.freq(1000)
led2.freq(1000)
led3.freq(1000)

DIM = 3000
OFF = 0

# OLED and I2C Init
i2c = I2C(1, sda=Pin(14), scl=Pin(15), freq=200000)
oled = SSD1306_I2C(WIDTH, HEIGHT, i2c)

buttons = [button1, button2, button3]
prev_states = [1, 1, 1]

def dot(x, y, r=2):
    oled.fill_rect(x - r, y - r, r * 2 + 1, r * 2 + 1, 1)

def draw_dice(value):
    x = 44
    y = 16
    size = 32

    oled.rect(x, y, size, size, 1)

    left = x + 8
    center_x = x + 16
    right = x + 24

    top = y + 8
    center_y = y + 16
    bottom = y + 24

    if value == 1:
        dot(center_x, center_y)

    elif value == 2:
        dot(left, top)
        dot(right, bottom)

    elif value == 3:
        dot(left, top)
        dot(center_x, center_y)
        dot(right, bottom)

    elif value == 4:
        dot(left, top)
        dot(right, top)
        dot(left, bottom)
        dot(right, bottom)

    elif value == 5:
        dot(left, top)
        dot(right, top)
        dot(center_x, center_y)
        dot(left, bottom)
        dot(right, bottom)

    elif value == 6:
        dot(left, top)
        dot(right, top)
        dot(left, center_y)
        dot(right, center_y)
        dot(left, bottom)
        dot(right, bottom)

# Start screen
oled.fill(0)
oled.text("Dice Roller", 20, 4)
oled.text("Press a button!", 8, 54)
draw_dice(1)
oled.show()

while True:
    for i in range(3):
        current_state = buttons[i].value()

        if prev_states[i] == 1 and current_state == 0:
            time.sleep_ms(50)

            if buttons[i].value() == 0:
                for step in range(8):
                    rolling_num = random.randint(1, 6)

                    led1.duty_u16(DIM if step % 3 == 0 else OFF)
                    led2.duty_u16(DIM if step % 3 == 1 else OFF)
                    led3.duty_u16(DIM if step % 3 == 2 else OFF)

                    oled.fill(0)
                    oled.text("Button " + str(i + 1), 31, 2)
                    oled.text("Rolling...", 31, 54)
                    draw_dice(rolling_num)
                    oled.show()
                    time.sleep_ms(120)

                result = random.randint(1, 6)

                print("Result:", result)

                oled.fill(0)
                oled.text("Button " + str(i + 1), 31, 2)
                oled.text("" + str(result), 58, 54)
                draw_dice(result)
                oled.show()

                led1.duty_u16(DIM)
                led2.duty_u16(DIM)
                led3.duty_u16(DIM)
                time.sleep_ms(1000)

                led1.duty_u16(OFF)
                led2.duty_u16(OFF)
                led3.duty_u16(OFF)

        prev_states[i] = current_state

    time.sleep_ms(10)