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"]


@dataclass
class AppState:
    stopEvent = Event()
    diskEventQueue: Queue[DiskEvent] = Queue()
    diskState: DiskState = "disconnected"


def onButtonPressed(state: AppState) -> None:
    state.diskState = "pairing"
    print("Changing disk state")


def diskPollWorker(state) -> 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":
            state.diskEventQueue.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":
            state.diskEventQueue.put(
                DiskEvent(
                    event="disconnected",
                    node=f"{device.device_node}",
                    serial=f"{device.get('ID_SERIAL_SHORT', 'Unknown')}",
                )
            )

        if state.stopEvent.is_set():
            return


def ledWorker(state: AppState) -> None:
    yellowCtr = 0
    previousState: DiskState | None = None

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

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

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

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

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

        yellowCtr += 1

        state.stopEvent.wait(0.25)


def main() -> None:
    try:
        state = AppState()
        button.when_activated = partial(onButtonPressed, state)

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

        diskPollThread.start()

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

        ledThread.start()

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

    finally:
        state.stopEvent.set()
        return


if __name__ == "__main__":
    main()