#!/usr/bin/env python3
"""UART finger controller for Radxa / Dragon Q6A (sensor-side receiver).

Combines ORY-FINGER processing (same path as ``examples/finger_pipeline.py``)
with a line-based UART command interface. A host (or ``uart_send.py``-style
transmitter) sends commands; this script replies and streams forces.

Dependencies (beyond orisys):
  pip install pyserial

Default UART (Dragon Q6A):
  port=/dev/ttyHS1  baud=115200  8N1

Commands (one line, case-insensitive):
  search              List cameras (/dev/video*) via orisys.capture.list_cameras
  connect [source]    Open FingerRuntime (camera index, /dev/videoN, or video file)
  run                 Stream FNORMAL / FSHEARX / FSHEARY until stop (or EOF)
  stop                Stop force streaming; keep source connected
  quit | exit         Close runtime and serial, then exit

Replies:
  camera,<index>,<path>,<WxH>
  OK search <N>
  OK connect <source>
  OK run
  force,<fnormal>,<fshearx>,<fsheary>
  OK eof              Video file ended (when --no-loop-video)
  OK stop
  OK quit
  ERR <context> <message>

Examples:
  # Live camera over UART
  python scripts/example_uart.py --port /dev/ttyHS1 --config finger

  # Offline debug from a recorded file (no camera needed)
  python scripts/example_uart.py --port /dev/ttyHS1 --video /path/to/clip.mp4
  # then UART: connect   (uses --video)  or  connect /path/to/clip.mp4
  # then UART: run
"""

from __future__ import annotations

import argparse
import os
import shlex
import sys
from pathlib import Path
from typing import Optional, Union

os.environ.setdefault("OPENCV_VIDEOIO_MSMF_ENABLE_HW_TRANSFORMS", "0")

try:
    import serial
except ImportError as exc:  # pragma: no cover
    raise SystemExit(
        "pyserial is required for this script.\n"
        "  pip install pyserial\n"
        f"Import error: {exc}"
    ) from exc

import cv2
from orisys.capture import is_live_camera, list_cameras
from orisys.finger import FingerRuntime

Source = Union[int, str]


def _parse_source(token: str) -> Optional[Source]:
    """Accept camera index int, numeric string, V4L path, or filesystem path."""
    text = str(token).strip().strip('"').strip("'")
    if not text:
        return None
    if text.isdigit() or (text.startswith("-") and text[1:].isdigit()):
        return int(text)
    return text


def _is_video_file(source: Source) -> bool:
    """True when source is a filesystem video path (not a live camera)."""
    if is_live_camera(source):
        return False
    if isinstance(source, int):
        return False
    path = Path(str(source))
    return path.is_file() or path.suffix.lower() in {
        ".mp4",
        ".avi",
        ".mkv",
        ".mov",
        ".m4v",
        ".wmv",
        ".webm",
    }


class UartFingerController:
    """Serial command dispatcher around a single FingerRuntime."""

    def __init__(
        self,
        ser: serial.Serial,
        *,
        config: str = "finger",
        default_source: Source = 0,
        max_camera_index: int = 10,
        loop_video: bool = True,
        show: bool = False,
        verbose: bool = False,
    ):
        self.ser = ser
        self.config = config
        self.default_source = default_source
        self.max_camera_index = max_camera_index
        self.loop_video = loop_video
        self.show = show
        self.verbose = verbose

        self.runtime: Optional[FingerRuntime] = None
        self._connected_source: Optional[Source] = None
        self._last_search_source: Optional[Source] = None
        self._streaming = False
        self._should_quit = False

    def reply(self, line: str) -> None:
        text = line.rstrip("\r\n") + "\n"
        self.ser.write(text.encode("utf-8"))
        self.ser.flush()
        if self.verbose:
            print(f"<< {text.rstrip()}", flush=True)

    def close_runtime(self) -> None:
        self._streaming = False
        if self.runtime is not None:
            try:
                self.runtime.close()
            except Exception:  # noqa: BLE001
                pass
            self.runtime = None
            self._connected_source = None
        if self.show:
            try:
                cv2.destroyAllWindows()
            except Exception:  # noqa: BLE001
                pass

    def _resolve_connect_source(self, args: list[str]) -> Source:
        if args:
            # Join remaining tokens so unquoted paths with spaces still work.
            joined = " ".join(args).strip()
            parsed = _parse_source(joined)
            if parsed is None:
                raise ValueError("empty connect source")
            return parsed
        # Prefer --video default over last search so offline debug stays on file.
        if isinstance(self.default_source, str) and not is_live_camera(self.default_source):
            return self.default_source
        if self._last_search_source is not None:
            return self._last_search_source
        return self.default_source

    def cmd_search(self, _args: list[str]) -> None:
        try:
            cameras = list_cameras(max_index=self.max_camera_index)
        except Exception as exc:  # noqa: BLE001
            self.reply(f"ERR search {exc}")
            return

        self._last_search_source = None
        for cam in cameras:
            path = (cam.path or f"/dev/video{cam.index}").strip() or f"/dev/video{cam.index}"
            size = f"{cam.width}x{cam.height}"
            self.reply(f"camera,{cam.index},{path},{size}")
            if self._last_search_source is None:
                self._last_search_source = cam.index
        self.reply(f"OK search {len(cameras)}")

    def cmd_connect(self, args: list[str]) -> None:
        try:
            source = self._resolve_connect_source(args)
        except ValueError as exc:
            self.reply(f"ERR connect {exc}")
            return

        if isinstance(source, str) and not is_live_camera(source):
            path = Path(source)
            if not path.is_file():
                self.reply(f"ERR connect file not found: {source}")
                return

        self._streaming = False
        self.close_runtime()
        try:
            runtime = FingerRuntime.open(
                config=self.config,
                video=source,
                verbose=self.verbose,
            )
            runtime.warmup(color=False)
        except Exception as exc:  # noqa: BLE001
            self.reply(f"ERR connect {exc}")
            return

        self.runtime = runtime
        self._connected_source = source
        kind = "file" if _is_video_file(source) else "camera"
        self.reply(f"OK connect {source} ({kind})")

    def _rewind_video(self) -> bool:
        """Seek connected file source back to frame 0. Return True on success."""
        if self.runtime is None or self._connected_source is None:
            return False
        if not _is_video_file(self._connected_source):
            return False
        cap = self.runtime.capture
        if cap is None:
            return False
        try:
            return bool(cap.set(cv2.CAP_PROP_POS_FRAMES, 0))
        except Exception:  # noqa: BLE001
            return False

    def _poll_command_line(self, *, blocking: bool = True) -> Optional[str]:
        """Read one command line if available; return None on timeout/empty."""
        if not blocking and self.ser.in_waiting <= 0:
            return None
        raw = self.ser.readline()
        if not raw:
            return None
        try:
            line = raw.decode("utf-8", errors="replace").strip()
        except Exception:  # noqa: BLE001
            return None
        return line or None

    def cmd_run(self, _args: list[str]) -> None:
        if self.runtime is None:
            self.reply("ERR run not connected")
            return

        self._streaming = True
        self.reply("OK run")
        tactile = self.runtime.tactile
        is_file = self._connected_source is not None and _is_video_file(self._connected_source)

        while self._streaming and not self._should_quit:
            ok, unrolled = self.runtime.grab_and_process(color=False)
            if not ok:
                if is_file and self.loop_video and self._rewind_video():
                    if self.verbose:
                        print("video EOF: looping", flush=True)
                    continue
                if is_file:
                    self.reply("OK eof")
                else:
                    self.reply("ERR run grab failed")
                self._streaming = False
                break

            tactile.compute_deformation()
            tactile.compute_contact()
            fnormal, fshearx, fsheary = tactile.read_info(
                tactile.info.FNORMAL,
                tactile.info.FSHEARX,
                tactile.info.FSHEARY,
            )
            self.reply(f"force,{float(fnormal):.6f},{float(fshearx):.6f},{float(fsheary):.6f}")

            if self.show:
                try:
                    import orisys

                    arrows = orisys.util.draw_arrows(
                        unrolled,
                        tactile.flow,
                        threshold=0.4,
                        grid_spacing=10,
                        arrow_scale=1.0,
                        valid_mask=tactile.valid_mask,
                    )
                    cv2.imshow("UV Unrolled", arrows)
                    key = cv2.waitKey(1) & 0xFF
                    if key == ord("q"):
                        self._should_quit = True
                        break
                    if key == ord("r"):
                        tactile.reset()
                except Exception as exc:  # noqa: BLE001
                    if self.verbose:
                        print(f"show error: {exc}", file=sys.stderr)

            # Non-blocking command poll so force streaming is not delayed by UART timeout.
            pending = self._poll_command_line(blocking=False)
            if pending is not None:
                if self.verbose:
                    print(f">> {pending}", flush=True)
                if not self._dispatch(pending, allow_run=False):
                    break

    def cmd_stop(self, _args: list[str]) -> None:
        self._streaming = False
        self.reply("OK stop")

    def cmd_quit(self, _args: list[str]) -> None:
        self._streaming = False
        self._should_quit = True
        self.reply("OK quit")

    def _dispatch(self, line: str, *, allow_run: bool = True) -> bool:
        """Dispatch one command line. Return False if the main loop should exit."""
        try:
            parts = shlex.split(line)
        except ValueError:
            parts = line.split()
        if not parts:
            return True
        cmd = parts[0].lower()
        args = parts[1:]

        if cmd == "search":
            self.cmd_search(args)
        elif cmd == "connect":
            self.cmd_connect(args)
        elif cmd == "run":
            if not allow_run:
                self.reply("ERR run already running")
            else:
                self.cmd_run(args)
        elif cmd == "stop":
            self.cmd_stop(args)
        elif cmd in ("quit", "exit"):
            self.cmd_quit(args)
            return False
        else:
            self.reply(f"ERR unknown {cmd}")
        return not self._should_quit

    def serve(self) -> None:
        self.reply("OK ready")
        while not self._should_quit:
            line = self._poll_command_line()
            if line is None:
                continue
            if self.verbose:
                print(f">> {line}", flush=True)
            if not self._dispatch(line):
                break


def open_serial(port: str, baud: int, timeout: float) -> serial.Serial:
    try:
        return serial.Serial(
            port=port,
            baudrate=baud,
            bytesize=serial.EIGHTBITS,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            timeout=timeout,
        )
    except serial.SerialException as exc:
        print(f"Serial error: {exc}", file=sys.stderr)
        print("Check:", file=sys.stderr)
        print(f"  1. Device node {port} exists", file=sys.stderr)
        print("  2. User has access (sudo or dialout group)", file=sys.stderr)
        print("  3. Wiring / USB-TTL connection", file=sys.stderr)
        raise SystemExit(1) from exc


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="UART finger controller (search/connect/run/stop force stream)"
    )
    parser.add_argument("--port", default="/dev/ttyHS1", help="Serial device (default: /dev/ttyHS1)")
    parser.add_argument("--baud", type=int, default=115200, help="Baud rate (default: 115200)")
    parser.add_argument(
        "--timeout",
        type=float,
        default=0.1,
        help="Serial readline timeout seconds (default: 0.1)",
    )
    parser.add_argument("--config", "-c", default="finger", help="Orisys config name or path")
    parser.add_argument(
        "--camera-id",
        type=int,
        default=0,
        help="Default camera index when connect has no argument (ignored if --video is set)",
    )
    parser.add_argument(
        "--video",
        default=None,
        help="Default video file for connect (offline debug without a camera)",
    )
    parser.add_argument(
        "--max-camera-index",
        type=int,
        default=10,
        help="Highest index probed by search",
    )
    parser.add_argument(
        "--loop-video",
        dest="loop_video",
        action="store_true",
        default=True,
        help="Loop video file on EOF during run (default: on)",
    )
    parser.add_argument(
        "--no-loop-video",
        dest="loop_video",
        action="store_false",
        help="Stop run with OK eof when video file ends",
    )
    parser.add_argument(
        "--show",
        action="store_true",
        help="Show UV unrolled OpenCV window (local debug only)",
    )
    parser.add_argument("--verbose", action="store_true", help="Echo UART traffic and SDK logs")
    args = parser.parse_args(argv)

    if args.video:
        video_path = Path(args.video)
        if not video_path.is_file():
            print(f"Video file not found: {args.video}", file=sys.stderr)
            return 1
        default_source: Source = str(video_path.resolve())
    else:
        default_source = args.camera_id

    ser = open_serial(args.port, args.baud, args.timeout)
    controller = UartFingerController(
        ser,
        config=args.config,
        default_source=default_source,
        max_camera_index=args.max_camera_index,
        loop_video=args.loop_video,
        show=args.show,
        verbose=args.verbose,
    )

    print(f"UART finger controller on {args.port} @ {args.baud} bps", flush=True)
    print(f"Default source: {default_source}", flush=True)
    print("Commands: search | connect [source|file] | run | stop | quit", flush=True)

    try:
        controller.serve()
    except KeyboardInterrupt:
        print("\nInterrupted.", flush=True)
        try:
            controller.reply("OK quit")
        except Exception:  # noqa: BLE001
            pass
    finally:
        controller.close_runtime()
        if ser.is_open:
            ser.close()
            print("Serial closed.", flush=True)
    return 0


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