from __future__ import annotations

import argparse
import csv
import struct
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import BinaryIO


BAUD_RATE = 921_600
UINT32_MASK = 0xFFFFFFFF

TURN_MAGIC = b"\xE7\x54\xC3\xA1"
PULSE16_MAGIC = b"\xE7\x50\xC3\xA1"
PULSE32_MAGIC = b"\xE7\x70\xC3\xA1"
GP_MAGIC = b"\xE7\x47\xC3\xA1"
MAGICS = {
    TURN_MAGIC: "turn",
    PULSE16_MAGIC: "pulse16",
    PULSE32_MAGIC: "pulse32",
    GP_MAGIC: "gp",
}

GP_COLUMNS = ("gp0_falling", "gp1_falling")
CSV_COLUMNS = ("time_us", "turn", "pulse", *GP_COLUMNS)


class EventTable:
    def __init__(self) -> None:
        self.offset: int | None = None
        self.rows: dict[int, list[int]] = {}

    def add(self, timestamp: int, column: str) -> None:
        if self.offset is None:
            self.offset = timestamp

        time_us = (timestamp - self.offset) & UINT32_MASK
        row = self.rows.setdefault(time_us, [0] * (len(CSV_COLUMNS) - 1))
        row[CSV_COLUMNS.index(column) - 1] = 1

    def write_csv(self, path: Path) -> None:
        path.parent.mkdir(parents=True, exist_ok=True)
        with path.open("w", newline="") as file:
            writer = csv.writer(file)
            writer.writerow(CSV_COLUMNS)
            for time_us in sorted(self.rows):
                writer.writerow((time_us, *self.rows[time_us]))


def choose_port() -> str | None:
    from serial.tools import list_ports

    ports = list(list_ports.comports())
    if not ports:
        print("No serial ports found.")
        return None

    print("Available serial ports:")
    for index, port in enumerate(ports, start=1):
        description = port.description or "serial port"
        print(f"  {index}. {port.device} - {description}")

    while True:
        choice = input("Choose a port number, or q to quit: ").strip().lower()
        if choice == "q":
            return None

        try:
            index = int(choice)
        except ValueError:
            print("Please enter a port number.")
            continue

        if 1 <= index <= len(ports):
            return ports[index - 1].device

        print(f"Please choose a number between 1 and {len(ports)}.")


def deadline_expired(deadline: float | None) -> bool:
    return deadline is not None and time.monotonic() >= deadline


def read_exact(
    stream: BinaryIO,
    size: int,
    *,
    deadline: float | None,
    eof_on_empty: bool,
) -> bytes | None:
    data = bytearray()
    while len(data) < size:
        chunk = stream.read(size - len(data))
        if not chunk:
            if eof_on_empty or deadline_expired(deadline):
                return None
            continue
        data.extend(chunk)
    return bytes(data)


def read_magic(
    stream: BinaryIO,
    *,
    deadline: float | None,
    eof_on_empty: bool,
) -> str | None:
    window = bytearray()

    while True:
        byte = stream.read(1)
        if not byte:
            if eof_on_empty or deadline_expired(deadline):
                return None
            continue

        window.extend(byte)
        if len(window) > 4:
            del window[0]

        if len(window) == 4:
            frame_type = MAGICS.get(bytes(window))
            if frame_type is not None:
                return frame_type


def parse_stream(
    stream: BinaryIO,
    events: EventTable,
    duration_s: float | None = None,
    eof_on_empty: bool = True,
) -> int:
    previous_pulse_timestamp: int | None = None
    deadline = None if duration_s is None else time.monotonic() + duration_s
    frames = 0

    while deadline is None or time.monotonic() < deadline:
        frame_type = read_magic(
            stream,
            deadline=deadline,
            eof_on_empty=eof_on_empty,
        )
        if frame_type is None:
            break

        if frame_type == "turn":
            payload = read_exact(
                stream,
                4,
                deadline=deadline,
                eof_on_empty=eof_on_empty,
            )
            if payload is None:
                break
            timestamp = struct.unpack("<I", payload)[0]
            previous_pulse_timestamp = timestamp
            events.add(timestamp, "turn")
        elif frame_type == "pulse16":
            payload = read_exact(
                stream,
                16 * 2,
                deadline=deadline,
                eof_on_empty=eof_on_empty,
            )
            if payload is None:
                break
            if previous_pulse_timestamp is not None:
                for delta in struct.unpack("<16H", payload):
                    if delta == 0:
                        continue
                    previous_pulse_timestamp = (
                        previous_pulse_timestamp + delta
                    ) & UINT32_MASK
                    events.add(previous_pulse_timestamp, "pulse")
        elif frame_type == "pulse32":
            payload = read_exact(
                stream,
                16 * 4,
                deadline=deadline,
                eof_on_empty=eof_on_empty,
            )
            if payload is None:
                break
            if previous_pulse_timestamp is not None:
                for delta in struct.unpack("<16I", payload):
                    if delta == 0:
                        continue
                    previous_pulse_timestamp = (
                        previous_pulse_timestamp + delta
                    ) & UINT32_MASK
                    events.add(previous_pulse_timestamp, "pulse")
        elif frame_type == "gp":
            payload = read_exact(
                stream,
                5,
                deadline=deadline,
                eof_on_empty=eof_on_empty,
            )
            if payload is None:
                break
            channel = payload[0]
            timestamp = struct.unpack("<I", payload[1:])[0]
            if channel < len(GP_COLUMNS):
                events.add(timestamp, GP_COLUMNS[channel])

        frames += 1

    return frames


def default_output_path() -> Path:
    started_at = datetime.now().strftime("%Y%m%d_%H%M%S")
    return Path("recordings") / f"{started_at}_events.csv"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Parse the ESP32 encoder binary stream into an event CSV."
    )
    source = parser.add_mutually_exclusive_group()
    source.add_argument("--port", help="Serial port to read from.")
    source.add_argument("--input", type=Path, help="Captured binary stream to parse.")
    parser.add_argument(
        "-o",
        "--output",
        type=Path,
        default=default_output_path(),
        help="Output CSV path.",
    )
    parser.add_argument("--baud", type=int, default=BAUD_RATE)
    parser.add_argument(
        "--duration",
        type=float,
        help="Recording duration in seconds. Without this, serial reads until Ctrl-C.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    events = EventTable()

    try:
        if args.input is not None:
            with args.input.open("rb") as stream:
                frames = parse_stream(stream, events)
        else:
            import serial

            port = args.port or choose_port()
            if port is None:
                return 0

            try:
                with serial.Serial(port, args.baud, timeout=0.1) as stream:
                    stream.reset_input_buffer()
                    print(f"Reading {port} at {args.baud} baud. Press Ctrl-C to stop.")
                    frames = parse_stream(
                        stream,
                        events,
                        args.duration,
                        eof_on_empty=False,
                    )
            except serial.SerialException as exc:
                print(f"Serial error: {exc}", file=sys.stderr)
                return 1
    except KeyboardInterrupt:
        print()
        frames = len(events.rows)
    except ImportError as exc:
        print(f"Missing Python dependency: {exc}", file=sys.stderr)
        return 1

    events.write_csv(args.output)
    print(f"Wrote {len(events.rows)} event times from {frames} frames to {args.output}")
    return 0


if __name__ == "__main__":
    sys.exit(main())