from datetime import datetime
from functools import partial
import json
from os import getuid
from pathlib import Path
from queue import Empty, Queue
import subprocess
from time import monotonic, sleep
from typing import Literal
import pyudev
from gpiozero import LED, Button
from threading import Thread, Event
from dataclasses import asdict, dataclass
import numpy as np
from libcamera import controls
from picamera2 import Picamera2
from picamera2.encoders import H264Encoder
from picamera2.outputs import CircularOutput

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"]
CameraState = Literal["off", "armed", "recording"]


@dataclass
class Config:
    diskSerial: str = ""


@dataclass
class AppState:
    stopEvent = Event()
    cameraStopComplete = Event()
    diskEventQueue: Queue[DiskEvent] = Queue()
    diskState: DiskState = "disconnected"
    config = Config()
    failedUmount: bool = False
    cameraState = "off"


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
    blueCtr = 0
    previousDiskState: DiskState | None = None
    previousCameraState: CameraState | None = None

    while not state.stopEvent.is_set():
        if state.diskState != previousDiskState:
            yellowCtr = 0
            previousDiskState = 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

        if state.cameraState != previousCameraState:
            blueCtr = 0
            previousCameraState = state.cameraState

        if state.cameraState == "off":
            blue.off()
        elif state.cameraState == "armed":
            blue.on()
        elif state.cameraState == "recording":
            if blueCtr % 4 < 2:
                blue.on()
            else:
                blue.off()

        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:
        Path(f"/mnt/{state.config.diskSerial}").mkdir(parents=True, exist_ok=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:
    if not state.failedUmount:
        try:
            subprocess.run(
                [
                    "umount",
                    f"/mnt/{state.config.diskSerial}",
                ],
                check=True,
            )
            print("unmounting disk")
        except subprocess.CalledProcessError:
            state.failedUmount = True
        else:
            state.diskState = "disconnected"
    else:
        print("force unmounting")
        subprocess.run(["fuser", "-km", f"/mnt/{state.config.diskSerial}"], check=False)
        subprocess.run(["sync"], check=False)
        subprocess.run(
            [
                "umount",
                "-f",
                f"/mnt/{state.config.diskSerial}",
            ],
            check=False,
        )
        state.diskState = "disconnected"
        state.failedUmount = False


def cameraWorker(state: AppState) -> None:
    VIDEO_SIZE = (1280, 720)
    ANALYSIS_SIZE = (320, 240)

    FRAME_RATE = 15
    BITRATE = 2_000_000

    QUIET_SECONDS = 10.0

    PRE_RECORD_SECONDS = 2.0

    WARMUP_SECONDS = 3.0

    PIXEL_CHANGE_THRESHOLD = 25

    MOTION_AREA_THRESHOLD = 0.015

    ANALYSIS_SAMPLE_STEP = 2

    OUTPUT_DIRECTORY = Path(f"/mnt/{state.config.diskSerial}/cctv")

    camera = Picamera2()

    configuration = camera.create_video_configuration(
        main={
            "size": VIDEO_SIZE,
            "format": "YUV420",
        },
        lores={
            "size": ANALYSIS_SIZE,
            "format": "YUV420",
        },
        controls={
            "FrameRate": FRAME_RATE,
        },
        buffer_count=6,
    )

    camera.configure(configuration)

    camera.set_controls(
        {
            "AfMode": controls.AfModeEnum.Continuous,
        }
    )

    encoder = H264Encoder(
        bitrate=BITRATE,
        repeat=True,
        iperiod=FRAME_RATE,
    )

    prebuffer_frames = max(
        1,
        round(FRAME_RATE * PRE_RECORD_SECONDS),
    )

    circular_output = CircularOutput(
        buffersize=prebuffer_frames,
    )

    encoder.output = circular_output

    recording = False
    recording_path: Path | None = None
    recording_trigger_time = 0.0
    last_motion_time = 0.0

    camera_started = False
    encoder_started = False

    try:
        camera.start()
        camera_started = True

        camera.start_encoder(encoder)
        encoder_started = True

        sleep(WARMUP_SECONDS)

        analysis_width, analysis_height = ANALYSIS_SIZE

        first_frame = camera.capture_array("lores")
        previous_frame = first_frame[:analysis_height, :analysis_width]
        previous_frame = previous_frame[
            ::ANALYSIS_SAMPLE_STEP,
            ::ANALYSIS_SAMPLE_STEP,
        ].astype(np.int16)

        while not state.stopEvent.is_set():
            if state.diskState != "connected":
                state.cameraState = "off"
                state.stopEvent.wait(1)
                continue
            if state.cameraState == "off":
                state.cameraState = "armed"
            frame = camera.capture_array("lores")

            # YUV420 starts with the full-resolution greyscale Y plane.
            current_frame = frame[:analysis_height, :analysis_width]
            current_frame = current_frame[
                ::ANALYSIS_SAMPLE_STEP,
                ::ANALYSIS_SAMPLE_STEP,
            ].astype(np.int16)

            difference = np.abs(current_frame - previous_frame)

            changed_pixels = np.count_nonzero(difference >= PIXEL_CHANGE_THRESHOLD)

            changed_fraction = changed_pixels / difference.size
            motion_detected = changed_fraction >= MOTION_AREA_THRESHOLD

            now = monotonic()

            if motion_detected:
                last_motion_time = now

                if not recording:
                    if not OUTPUT_DIRECTORY.exists():
                        OUTPUT_DIRECTORY.mkdir(parents=True, exist_ok=True)

                    filename_timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")

                    recording_path = OUTPUT_DIRECTORY / f"{filename_timestamp}.h264"

                    # Start writing the circular pre-buffer and all subsequent
                    # encoded H.264 packets directly to the file.
                    circular_output.fileoutput = str(recording_path)
                    circular_output.start()

                    recording = True
                    state.cameraState = "recording"
                    recording_trigger_time = now

            elif recording and now - last_motion_time >= QUIET_SECONDS:
                circular_output.stop()
                state.cameraState = "armed"
                recording = False

                elapsed = now - recording_trigger_time

                recording_path = None

            previous_frame = current_frame
    finally:
        if recording:
            circular_output.stop()

        if encoder_started:
            try:
                camera.stop_encoder()
            except Exception:
                pass

        if camera_started:
            try:
                camera.stop()
            except Exception:
                pass

        camera.close()
        state.cameraStopComplete.set()


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

        cameraThread = Thread(
            target=cameraWorker,
            name="camera_thread",
            args=(state,),
            daemon=True,
        )

        cameraThread.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()
        state.cameraStopComplete.wait()
        if state.diskState == "connected":
            unmountDisk(state)
            if state.failedUmount:
                unmountDisk(state)
        return


if __name__ == "__main__":
    main()