from functools import partial
import json
from os import getuid
from pathlib import Path
from queue import Empty, Queue
import subprocess
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:
    if state.diskState == "disconnected":
        state.diskState = "pairing"
    elif state.diskState == "pairing":
        state.diskState = "disconnected"
    elif state.diskState == "connected":
        unmountDisk(state)
    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 mountDisk(state: AppState) -> None:
    state.diskState = "connecting"
    node = findDiskNode(state.config.diskSerial)
    try:
        subprocess.run(
            ["mkdir", "-p", f"/mnt/{state.config.diskSerial}"],
            check=True,
        )
        sleep(1)
        subprocess.run(
            [
                "mount",
                f"{node}1",
                f"/mnt/{state.config.diskSerial}",
                "-o",
                f"uid={getuid()}",
            ],
            check=True,
        )
        sleep(1)
    except subprocess.CalledProcessError:
        pass
    else:
        state.diskState = "connected"


def unmountDisk(state: AppState) -> None:
    try:
        subprocess.run(
            [
                "umount",
                f"/mnt/{state.config.diskSerial}",
            ],
            check=True,
        )
        print("unmounting disk")
    except subprocess.CalledProcessError:
        pass
    else:
        state.diskState = "disconnected"


def main() -> None:
    try:
        state = AppState()
        state.config = loadOrCreateConfig(Path("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()

        if state.config.diskSerial != "":
            mountDisk(state)

        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, Path("config.json"))
                    print(findDiskNode(state.config.diskSerial))
                    mountDisk(state)

                if (
                    state.diskState == "disconnected"
                    and event.event == "connected"
                    and event.serial == state.config.diskSerial
                ):
                    mountDisk(state)

            except Empty:
                sleep(1)

    finally:
        state.stopEvent.set()
        if state.diskState == "connected":
            unmountDisk(state)
        return


if __name__ == "__main__":
    main()