from functools import partial
import json
from pathlib import Path
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 asdict, 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 Config:
    diskSerial: str = ""


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


def saveConfig(config: Config, path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)

    with path.open("w", encoding="utf-8") as file:
        json.dump(asdict(config), file, indent=4)


def loadOrCreateConfig(path: Path) -> Config:
    if not path.exists():
        config = Config()
        saveConfig(config, path)
        return config

    with path.open("r", encoding="utf-8") as file:
        rawConfig = json.load(file)

    return Config(**rawConfig)


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 findDiskNode(serial: str) -> str | None:
    context = pyudev.Context()

    for device in context.list_devices(
        subsystem="block",
        DEVTYPE="disk",
    ):
        deviceSerial = device.get("ID_SERIAL_SHORT") or device.get("ID_SERIAL")

        if deviceSerial == serial:
            return device.device_node

    return None


def main() -> None:
    try:
        state = AppState()
        state.config = loadOrCreateConfig("config.json")
        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()
                if (
                    state.diskState == "pairing"
                    and event.event == "connected"
                    and event.serial != "Unknown"
                ):
                    state.config.diskSerial = event.serial
                    saveConfig(state.config, "config.json")
                    print(findDiskNode(state.config.diskSerial))
                    state.diskState = "connected"

            except Empty:
                sleep(1)

    finally:
        state.stopEvent.set()
        return


if __name__ == "__main__":
    main()