from machine import Pin, I2C
import time
import ssd1306
import ujson

SETTINGS_FILE = "settings.json"

I2C_SDA = 0
I2C_SCL = 1

BTN_PLUS_PIN  = 2
BTN_MINUS_PIN = 3
BTN_DIGIT_PIN = 4
BTN_START_PIN = 5

SOL1_PIN = 6
SOL2_PIN = 7

MIN_TIME_MS = 50
MAX_TIME_MS = 10800000
DEFAULT_SETPOINT_MS = 10000

PULSE_MS = 80
POST_PULSE_FINISH_MS = 120

DEBOUNCE_MS = 40
SHORT_REPEAT_START_MS = 450
FAST_REPEAT_START_MS = 1800
SHORT_REPEAT_INTERVAL_MS = 180
FAST_REPEAT_INTERVAL_MS = 90

BLINK_INTERVAL_MS = 500
SAVE_MIN_INTERVAL_MS = 400
DISPLAY_ZERO_THRESHOLD_MS = 100
DISPLAY_REFRESH_MS = 50

DIGIT_STEPS_MS = [
    10000000,  # 10000 s
    1000000,   # 1000 s
    100000,    # 100 s
    10000,     # 10 s
    1000,      # 1 s
    100,       # 100 ms
    10,        # 10 ms
    1          # 1 ms
]

i2c = I2C(0, sda=Pin(I2C_SDA), scl=Pin(I2C_SCL), freq=400000)
oled = ssd1306.SSD1306_I2C(128, 64, i2c)

btn_plus  = Pin(BTN_PLUS_PIN, Pin.IN, Pin.PULL_UP)
btn_minus = Pin(BTN_MINUS_PIN, Pin.IN, Pin.PULL_UP)
btn_digit = Pin(BTN_DIGIT_PIN, Pin.IN, Pin.PULL_UP)
btn_start = Pin(BTN_START_PIN, Pin.IN, Pin.PULL_UP)

sol1 = Pin(SOL1_PIN, Pin.OUT)
sol2 = Pin(SOL2_PIN, Pin.OUT)

sol1.value(0)
sol2.value(0)

setpoint_ms = DEFAULT_SETPOINT_MS
selected_digit = 0
last_status = "Pronto"

blink_state = True
last_blink_ms = 0
last_save_ms = 0
last_display_refresh_ms = 0
screen_dirty = True

STATE_IDLE = 0
STATE_COUNTING = 1
STATE_FINISHING = 2

cycle_state = STATE_IDLE

run_start_us = 0
target_ms = 0.0
second_shot_done = False
finish_at_ms = 0

pulse_active = False
pulse_end_ms = 0

stopped_elapsed_ms = 0

SEGMENTS = {
    "0": (1, 1, 1, 1, 1, 1, 0),
    "1": (0, 1, 1, 0, 0, 0, 0),
    "2": (1, 1, 0, 1, 1, 0, 1),
    "3": (1, 1, 1, 1, 0, 0, 1),
    "4": (0, 1, 1, 0, 0, 1, 1),
    "5": (1, 0, 1, 1, 0, 1, 1),
    "6": (1, 0, 1, 1, 1, 1, 1),
    "7": (1, 1, 1, 0, 0, 0, 0),
    "8": (1, 1, 1, 1, 1, 1, 1),
    "9": (1, 1, 1, 1, 0, 1, 1),
}

DIGIT_W = 14
DIGIT_H = 28
SEG_T = 2
DIGIT_SPACING = 20
COMMA_SPACING = 8
DISPLAY_Y = 12


def clamp(v, vmin, vmax):
    return max(vmin, min(vmax, v))


def now_ms():
    return time.ticks_ms()


def ticks_diff_ms(a, b):
    return time.ticks_diff(a, b)


def ticks_reached(now_value, target_value):
    return time.ticks_diff(now_value, target_value) >= 0


def elapsed_us(start_us):
    return time.ticks_diff(time.ticks_us(), start_us)


def get_elapsed_ms():
    if cycle_state == STATE_IDLE:
        return 0.0

    val_ms = elapsed_us(run_start_us) / 1000.0

    if val_ms < 0:
        val_ms = 0.0

    if val_ms < DISPLAY_ZERO_THRESHOLD_MS:
        val_ms = 0.0

    return val_ms


def format_time_fixed(ms):
    ms = int(ms)
    sec = ms // 1000
    milli = ms % 1000
    return "{:05d},{:03d}".format(sec, milli)


def format_time_trimmed(ms):
    ms = int(ms)
    sec = ms // 1000
    milli = ms % 1000
    return "{:d},{:03d}".format(sec, milli)


def mark_dirty():
    global screen_dirty
    screen_dirty = True


def save_settings(force=False):
    global last_save_ms
    t = now_ms()

    if (not force) and ticks_diff_ms(t, last_save_ms) < SAVE_MIN_INTERVAL_MS:
        return

    data = {
        "setpoint_ms": int(clamp(setpoint_ms, MIN_TIME_MS, MAX_TIME_MS))
    }

    try:
        with open(SETTINGS_FILE, "w") as f:
            ujson.dump(data, f)
        last_save_ms = t
    except Exception:
        pass


def load_settings():
    global setpoint_ms
    try:
        with open(SETTINGS_FILE, "r") as f:
            data = ujson.load(f)
        setpoint_ms = int(clamp(int(data.get("setpoint_ms", DEFAULT_SETPOINT_MS)), MIN_TIME_MS, MAX_TIME_MS))
    except Exception:
        setpoint_ms = DEFAULT_SETPOINT_MS


def update_digit(delta):
    global setpoint_ms

    step = DIGIT_STEPS_MS[selected_digit]
    new_value = setpoint_ms + (delta * step)
    new_value = int(clamp(new_value, MIN_TIME_MS, MAX_TIME_MS))

    if new_value != setpoint_ms:
        setpoint_ms = new_value
        save_settings()
        mark_dirty()


class SmartButton:
    def __init__(self, pin):
        self.pin = pin
        self.prev_pressed = False
        self.press_start = 0
        self.last_repeat = 0
        self.last_edge = 0

    def raw_pressed(self):
        return self.pin.value() == 0

    def update(self):
        t = now_ms()
        pressed = self.raw_pressed()

        if pressed != self.prev_pressed:
            if ticks_diff_ms(t, self.last_edge) < DEBOUNCE_MS:
                return None
            self.last_edge = t
            self.prev_pressed = pressed

            if pressed:
                self.press_start = t
                self.last_repeat = t
                return "press"
            return "release"

        if not pressed:
            return None

        held = ticks_diff_ms(t, self.press_start)

        if held >= FAST_REPEAT_START_MS:
            if ticks_diff_ms(t, self.last_repeat) >= FAST_REPEAT_INTERVAL_MS:
                self.last_repeat = t
                return "repeat2"
        elif held >= SHORT_REPEAT_START_MS:
            if ticks_diff_ms(t, self.last_repeat) >= SHORT_REPEAT_INTERVAL_MS:
                self.last_repeat = t
                return "repeat1"

        return None


btn_plus_s  = SmartButton(btn_plus)
btn_minus_s = SmartButton(btn_minus)
btn_digit_s = SmartButton(btn_digit)
btn_start_s = SmartButton(btn_start)


def stop_all_outputs():
    global pulse_active
    sol1.value(0)
    sol2.value(0)
    pulse_active = False


def start_pulse():
    global pulse_active, pulse_end_ms
    sol1.value(1)
    sol2.value(1)
    pulse_active = True
    pulse_end_ms = time.ticks_add(now_ms(), PULSE_MS)
    mark_dirty()


def update_pulse():
    global pulse_active
    if pulse_active and ticks_reached(now_ms(), pulse_end_ms):
        stop_all_outputs()
        mark_dirty()


def start_cycle():
    global cycle_state, run_start_us, target_ms, second_shot_done
    global last_status, stopped_elapsed_ms

    last_status = "Rodando"
    stopped_elapsed_ms = 0
    start_pulse()
    run_start_us = time.ticks_us()
    target_ms = setpoint_ms
    second_shot_done = False
    cycle_state = STATE_COUNTING
    mark_dirty()


def abort_cycle():
    global cycle_state, second_shot_done, last_status, stopped_elapsed_ms

    stopped_elapsed_ms = get_elapsed_ms()
    stop_all_outputs()
    cycle_state = STATE_IDLE
    second_shot_done = False
    last_status = "Interrompido"
    mark_dirty()


def finish_cycle():
    global cycle_state, last_status
    cycle_state = STATE_IDLE
    last_status = "Fim"
    save_settings(force=True)
    mark_dirty()


def update_cycle():
    global cycle_state, second_shot_done, finish_at_ms

    update_pulse()

    if cycle_state == STATE_COUNTING:
        current_ms = get_elapsed_ms()

        if (not second_shot_done) and current_ms >= target_ms:
            start_pulse()
            second_shot_done = True
            finish_at_ms = time.ticks_add(now_ms(), POST_PULSE_FINISH_MS)
            cycle_state = STATE_FINISHING
            mark_dirty()
            return

        if current_ms > (setpoint_ms + 5000):
            finish_cycle()
            return

    elif cycle_state == STATE_FINISHING:
        if ticks_reached(now_ms(), finish_at_ms) and (not pulse_active):
            finish_cycle()


def update_blink():
    global blink_state, last_blink_ms
    t = now_ms()
    if ticks_diff_ms(t, last_blink_ms) >= BLINK_INTERVAL_MS:
        blink_state = not blink_state
        last_blink_ms = t
        if cycle_state == STATE_IDLE:
            mark_dirty()


def draw_hseg_sharp(x, y, w, t):
    oled.fill_rect(x + 2, y, w - 4, t, 1)
    oled.fill_rect(x + 1, y + 1, w - 2, max(1, t - 1), 1)


def draw_vseg_sharp(x, y, t, h):
    oled.fill_rect(x, y + 2, t, h - 4, 1)
    oled.fill_rect(x + 1, y + 1, max(1, t - 1), h - 2, 1)


def draw_digit_sharp(x, y, ch, on=True):
    if ch not in SEGMENTS:
        return

    a, b, c, d, e, f, g = SEGMENTS[ch]
    if not on:
        a = b = c = d = e = f = g = 0

    w = DIGIT_W
    h = DIGIT_H
    t = SEG_T
    vh = (h - 3 * t) // 2

    if a:
        draw_hseg_sharp(x + t, y, w, t)
    if b:
        draw_vseg_sharp(x + w + t, y + t, t, vh)
    if c:
        draw_vseg_sharp(x + w + t, y + 2 * t + vh, t, vh)
    if d:
        draw_hseg_sharp(x + t, y + 2 * t + 2 * vh, w, t)
    if e:
        draw_vseg_sharp(x, y + 2 * t + vh, t, vh)
    if f:
        draw_vseg_sharp(x, y + t, t, vh)
    if g:
        draw_hseg_sharp(x + t, y + t + vh, w, t)


def draw_comma_sharp(x, y):
    oled.fill_rect(x + 1, y + DIGIT_H - 5, 2, 3, 1)
    oled.fill_rect(x + 2, y + DIGIT_H - 2, 1, 2, 1)


def draw_time_sharp(text, y=DISPLAY_Y, blink_selected=False):
    total_w = 0
    for ch in text:
        total_w += COMMA_SPACING if ch == "," else DIGIT_SPACING
    total_w -= 2

    x = max(0, (128 - total_w) // 2)
    digit_idx = 0

    for ch in text:
        if ch.isdigit():
            on = True
            if blink_selected and digit_idx == selected_digit and not blink_state:
                on = False
            draw_digit_sharp(x, y, ch, on=on)
            x += DIGIT_SPACING
            digit_idx += 1
        elif ch == ",":
            draw_comma_sharp(x, y)
            x += COMMA_SPACING


def draw_idle_screen():
    oled.fill(0)
    draw_time_sharp(format_time_fixed(setpoint_ms), DISPLAY_Y, blink_selected=True)
    oled.show()


def draw_running_screen():
    oled.fill(0)
    elapsed_ms = get_elapsed_ms()
    draw_time_sharp(format_time_trimmed(elapsed_ms), DISPLAY_Y, blink_selected=False)
    oled.show()


def draw_finished_screen():
    oled.fill(0)
    draw_time_sharp(format_time_trimmed(setpoint_ms), DISPLAY_Y, blink_selected=False)
    oled.show()


def draw_aborted_screen():
    oled.fill(0)
    draw_time_sharp(format_time_trimmed(stopped_elapsed_ms), DISPLAY_Y, blink_selected=False)
    oled.show()


def draw_screen():
    global screen_dirty, last_display_refresh_ms

    t = now_ms()

    if cycle_state in (STATE_COUNTING, STATE_FINISHING):
        if ticks_diff_ms(t, last_display_refresh_ms) >= DISPLAY_REFRESH_MS:
            screen_dirty = True
            last_display_refresh_ms = t

    if not screen_dirty:
        return

    if cycle_state in (STATE_COUNTING, STATE_FINISHING):
        draw_running_screen()
    elif last_status == "Fim":
        draw_finished_screen()
    elif last_status == "Interrompido":
        draw_aborted_screen()
    else:
        draw_idle_screen()

    screen_dirty = False


def handle_idle_inputs():
    global selected_digit, last_status

    evt_plus = btn_plus_s.update()
    evt_minus = btn_minus_s.update()
    evt_digit = btn_digit_s.update()
    evt_start = btn_start_s.update()

    if evt_plus in ("press", "repeat1", "repeat2"):
        update_digit(+1)
        last_status = "Editando"
        mark_dirty()

    if evt_minus in ("press", "repeat1", "repeat2"):
        update_digit(-1)
        last_status = "Editando"
        mark_dirty()

    if evt_digit == "press":
        selected_digit = (selected_digit + 1) % 8
        last_status = "Digito"
        mark_dirty()

    if evt_start == "press":
        start_cycle()


def handle_running_inputs():
    evt_start = btn_start_s.update()
    if evt_start == "press":
        abort_cycle()


load_settings()
mark_dirty()

while True:
    update_blink()
    update_cycle()

    if cycle_state == STATE_IDLE:
        handle_idle_inputs()
    else:
        handle_running_inputs()

    draw_screen()
    time.sleep_ms(10)