#!/usr/bin/env python3
import argparse
import struct
import threading
import time
from collections import deque
from queue import Empty, Queue

try:
    import pygame
except ImportError as exc:
    raise SystemExit(
        "Missing dependency: pygame. Install with `pip install pygame pyserial`."
    ) from exc

try:
    import serial
except ImportError as exc:
    raise SystemExit(
        "Missing dependency: pyserial. Install with `pip install pygame pyserial`."
    ) from exc


RECORD_STRUCT = struct.Struct("<IIIII")
RECORD_FIELDS = ("ms", "ticks_r", "target_r", "ticks_l", "target_l")

BACKGROUND = (18, 20, 24)
GRID = (48, 54, 64)
AXIS = (140, 148, 160)
TEXT = (230, 235, 240)
ERROR = (255, 110, 110)

TRACE_DEFS = [
    {
        "name": "speed_r",
        "label": "SPEED_R",
        "color": (80, 220, 140),
        "unit": "ticks/s",
        "kind": "derived",
        "compute": lambda window: compute_speed(window, "ticks_r"),
    },
    {
        "name": "target_r",
        "label": "TARGET_R",
        "color": (220, 80, 140),
        "unit": "ticks/s",
        "kind": "field",
        "field": "target_r",
    },
    {
        "name": "speed_l",
        "label": "SPEED_L",
        "color": (80, 170, 255),
        "unit": "ticks/s",
        "kind": "derived",
        "compute": lambda window: compute_speed(window, "ticks_l"),
    },
    {
        "name": "target_l",
        "label": "TARGET_L",
        "color": (255, 190, 80),
        "unit": "ticks/s",
        "kind": "field",
        "field": "target_l",
    },
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Read binary motor samples from a serial port and draw traces live."
    )
    parser.add_argument("port", help="Serial port, for example /dev/ttyUSB0 or COM3")
    parser.add_argument(
        "-b",
        "--baudrate",
        type=int,
        default=115200,
        help="Serial baudrate (default: 115200)",
    )
    parser.add_argument(
        "--history-seconds",
        type=float,
        default=10.0,
        help="Initial visible history window in seconds (default: 10)",
    )
    parser.add_argument(
        "--avg-window",
        type=int,
        default=5,
        help="Sliding average window in samples for derived traces (default: 5)",
    )
    parser.add_argument(
        "--max-samples",
        type=int,
        default=30000,
        help="Maximum number of samples to keep in memory (default: 30000)",
    )
    parser.add_argument(
        "--width",
        type=int,
        default=1200,
        help="Window width in pixels (default: 1200)",
    )
    parser.add_argument(
        "--height",
        type=int,
        default=700,
        help="Window height in pixels (default: 700)",
    )
    return parser.parse_args()


def serial_reader(port: str, baudrate: int, output: Queue) -> None:
    try:
        with serial.Serial(port, baudrate=baudrate, timeout=1) as ser:
            while True:
                raw = ser.read(RECORD_STRUCT.size)
                if len(raw) != RECORD_STRUCT.size:
                    continue
                values = RECORD_STRUCT.unpack(raw)
                sample = dict(zip(RECORD_FIELDS, values))
                sample["wall_time"] = time.monotonic()
                output.put(sample)
    except serial.SerialException as exc:
        output.put({"error": f"Serial error: {exc}"})


def compute_speed(
    raw_window: list[dict[str, float | int]],
    tick_field: str,
) -> float | None:
    if len(raw_window) < 2:
        return None

    first = raw_window[0]
    last = raw_window[-1]
    dt_ms = int(last["ms"]) - int(first["ms"])
    if dt_ms <= 0:
        return None

    dt_ticks = int(last[tick_field]) - int(first[tick_field])
    return (1000.0 * dt_ticks) / dt_ms


def update_traces(
    raw_samples: deque[dict[str, float | int]],
    trace_samples: dict[str, deque[tuple[float, float, dict[str, float | int]]]],
    avg_window: int,
) -> None:
    window = list(raw_samples)[-max(avg_window, 2):]
    sample = raw_samples[-1]

    for trace in TRACE_DEFS:
        value = None
        if trace["kind"] == "field":
            value = float(sample[trace["field"]])
        elif trace["kind"] == "derived":
            value = trace["compute"](window)

        if value is not None:
            trace_samples[trace["name"]].append(
                (float(sample["wall_time"]), float(value), sample)
            )


def drain_queue(
    queue: Queue,
    raw_samples: deque[dict[str, float | int]],
    trace_samples: dict[str, deque[tuple[float, float, dict[str, float | int]]]],
    avg_window: int,
) -> str | None:
    error_message = None
    while True:
        try:
            item = queue.get_nowait()
        except Empty:
            return error_message

        if "error" in item:
            error_message = str(item["error"])
            continue

        raw_samples.append(item)
        update_traces(raw_samples, trace_samples, avg_window)


def draw_grid(
    surface: pygame.Surface,
    rect: pygame.Rect,
    font: pygame.font.Font,
    view_span: float,
    y_max: float,
) -> None:
    pygame.draw.rect(surface, GRID, rect, width=1)

    for fraction in (0.0, 0.25, 0.5, 0.75, 1.0):
        y = rect.top + round(fraction * rect.height)
        pygame.draw.line(surface, GRID, (rect.left, y), (rect.right, y), 1)
        value = (1.0 - fraction) * y_max
        text = font.render(f"{value:.1f}", True, AXIS)
        surface.blit(text, (10, y - text.get_height() // 2))

    for fraction in (0.0, 0.25, 0.5, 0.75, 1.0):
        x = rect.left + round(fraction * rect.width)
        pygame.draw.line(surface, GRID, (x, rect.top), (x, rect.bottom), 1)
        seconds_ago = view_span * (1.0 - fraction)
        label = f"-{seconds_ago:.1f}s" if seconds_ago > 0.05 else "now"
        text = font.render(label, True, AXIS)
        surface.blit(text, (x - text.get_width() // 2, rect.bottom + 8))


def draw_trace(
    surface: pygame.Surface,
    rect: pygame.Rect,
    samples: deque[tuple[float, float, dict[str, float | int]]],
    view_span: float,
    now: float,
    color: tuple[int, int, int],
    y_max: float,
) -> tuple[float | None, dict[str, float | int] | None]:
    visible_points = []
    latest_value = None
    latest_sample = None

    for timestamp, value, sample in samples:
        age = now - timestamp
        if age < 0 or age > view_span:
            continue
        x = rect.right - (age / view_span) * rect.width
        y = rect.bottom - (value / y_max) * rect.height
        visible_points.append((round(x), round(y)))
        latest_value = value
        latest_sample = sample

    if len(visible_points) >= 2:
        pygame.draw.lines(surface, color, False, visible_points, 2)
    elif len(visible_points) == 1:
        pygame.draw.circle(surface, color, visible_points[0], 2)

    return latest_value, latest_sample


def zoom(view_span: float, direction: int, initial_span: float) -> float:
    if direction > 0:
        return max(1.0, view_span / 1.2)
    return min(max(initial_span * 20.0, 60.0), view_span * 1.2)


def main() -> int:
    args = parse_args()
    raw_samples: deque[dict[str, float | int]] = deque(maxlen=args.max_samples)
    trace_samples = {
        trace["name"]: deque(maxlen=args.max_samples) for trace in TRACE_DEFS
    }
    queue: Queue = Queue()

    thread = threading.Thread(
        target=serial_reader, args=(args.port, args.baudrate, queue), daemon=True
    )
    thread.start()

    pygame.init()
    pygame.display.set_caption(f"Motor traces: {args.port} @ {args.baudrate}")
    screen = pygame.display.set_mode((args.width, args.height), pygame.RESIZABLE)
    font = pygame.font.SysFont("monospace", 18)
    small_font = pygame.font.SysFont("monospace", 14)
    clock = pygame.time.Clock()

    view_span = max(args.history_seconds, 1.0)
    error_message = None
    running = True

    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.MOUSEWHEEL:
                view_span = zoom(view_span, event.y, args.history_seconds)
            elif event.type == pygame.MOUSEBUTTONDOWN:
                if event.button == 4:
                    view_span = zoom(view_span, 1, args.history_seconds)
                elif event.button == 5:
                    view_span = zoom(view_span, -1, args.history_seconds)

        queued_error = drain_queue(queue, raw_samples, trace_samples, args.avg_window)
        if queued_error is not None:
            error_message = queued_error

        width, height = screen.get_size()
        plot_rect = pygame.Rect(70, 40, max(100, width - 100), max(100, height - 100))
        now = time.monotonic()

        y_max = 1.0
        for trace in TRACE_DEFS:
            for timestamp, value, _sample in trace_samples[trace["name"]]:
                age = now - timestamp
                if 0 <= age <= view_span and value > y_max:
                    y_max = value
        y_max = max(10.0, y_max * 1.2)

        screen.fill(BACKGROUND)
        draw_grid(screen, plot_rect, small_font, view_span, y_max)

        latest_values: dict[str, float] = {}
        latest_sample: dict[str, float | int] | None = None

        for trace in TRACE_DEFS:
            value, sample = draw_trace(
                screen,
                plot_rect,
                trace_samples[trace["name"]],
                view_span,
                now,
                trace["color"],
                y_max,
            )
            if value is not None:
                latest_values[trace["name"]] = value
            if sample is not None:
                latest_sample = sample

        header = (
            f"port={args.port} baud={args.baudrate} zoom={view_span:.1f}s "
            f"avg_window={args.avg_window}"
        )
        screen.blit(font.render(header, True, TEXT), (10, 8))
        screen.blit(
            small_font.render("mouse wheel: zoom horizontal axis", True, AXIS),
            (10, height - 24),
        )

        legend_x = 10
        for trace in TRACE_DEFS:
            pygame.draw.line(
                screen,
                trace["color"],
                (legend_x, 43),
                (legend_x + 24, 43),
                3,
            )
            screen.blit(small_font.render(trace["label"], True, TEXT), (legend_x + 30, 34))
            legend_x += 120

        if error_message:
            error_surface = font.render(error_message, True, ERROR)
            screen.blit(error_surface, (10, 56))

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()
    return 0


if __name__ == "__main__":
    raise SystemExit(main())