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

red = LED(17)
yellow = LED(27)
blue = LED(22)


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


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


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)

    yellow.off()


def main() -> None:
    try:
        stopEvent = Event()
        diskEventQueue: Queue[DiskEvent] = Queue()
        diskState: DiskState = "pairing"

        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()