from functools import partial
from queue import Empty, Queue
from time import sleep
from typing import Literal
import pyudev
from gpiozero import LED, Button
from threading import Thread, Event
from dataclasses import dataclass

red = LED(17)
yellow = LED(27)
blue = LED(22)
button = Button(23, pull_up=True, bounce_time=0.05)


@dataclass
class DiskEvent:
    event: str
    node: str
    serial: str


DiskState = Literal["disconnected", "pairing", "connecting", "connected"]


def onButtonPressed(diskState: DiskState) -> None:
    diskState = "pairing"


def diskPollWorker(stopEvent: Event, eventQueue: Queue) -> None:
    context = pyudev.Context()

    monitor = pyudev.Monitor.from_netlink(context)
    monitor.filter_by(subsystem="block")

    print("Waiting for disks...")

    for device in iter(monitor.poll, None):
        if device.action == "add" and device.device_type == "disk":
            eventQueue.put(
                DiskEvent(
                    event="connected",
                    node=f"{device.device_node}",
                    serial=f"{device.get('ID_SERIAL_SHORT', 'Unknown')}",
                )
            )

        elif device.action == "remove" and device.device_type == "disk":
            eventQueue.put(
                DiskEvent(
                    event="disconnected",
                    node=f"{device.device_node}",
                    serial=f"{device.get('ID_SERIAL_SHORT', 'Unknown')}",
                )
            )

        if stopEvent.is_set():
            return


def ledWorker(stopEvent: Event, diskState: DiskState) -> None:
    yellowCtr = 0
    previousState: DiskState | None = None

    while not stopEvent.is_set():
        if diskState != previousState:
            yellowCtr = 0
            previousState = diskState

        if diskState == "disconnected":
            yellow.off()

        elif diskState == "connected":
            yellow.on()

        elif diskState == "pairing":
            if yellowCtr % 2 == 0:
                yellow.on()
            else:
                yellow.off()

        elif diskState == "connecting":
            if yellowCtr % 4 < 2:
                yellow.on()
            else:
                yellow.off()

        yellowCtr += 1

        stopEvent.wait(0.25)


def main() -> None:
    try:
        stopEvent = Event()
        diskEventQueue: Queue[DiskEvent] = Queue()
        diskState: DiskState = "disconnected"
        button.when_activated = partial(onButtonPressed, diskState)

        diskPollThread = Thread(
            target=diskPollWorker,
            name="disk_poll_thread",
            args=(stopEvent, diskEventQueue),
            daemon=True,
        )

        diskPollThread.start()

        ledThread = Thread(
            target=ledWorker,
            name="led_thread",
            args=(stopEvent, diskState),
            daemon=True,
        )

        ledThread.start()

        while True:
            try:
                event = diskEventQueue.get_nowait()
                print(event)
            except Empty:
                sleep(1)

    finally:
        stopEvent.set()
        return


if __name__ == "__main__":
    main()