#!/usr/bin/env python3
"""Headless UART finger service for 1-4 sensors; see ../docs/uart.md.

The login console and data UART must be separate ports. Single-sensor force
records retain the original format; multi-sensor records add a 1-based ID.
Importing this module opens no devices; acquisition is owned by finger_api.
"""

from __future__ import annotations

import argparse
import logging
import math
import os
import shlex
import sys
from collections import deque
from pathlib import Path

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

from finger_api import FingerSensorBank, is_video_file, parse_source

logger = logging.getLogger(__name__)


class UartFingerController:
    """Single-threaded command service; the caller owns serial and final cleanup."""

    def __init__(
        self, ser, *, config="finger", default_source=0, sources=None,
        configs=None, max_camera_index=10, loop_video=True, show=False,
        verbose=False, read_only=False, autostart=False,
    ):
        self.ser = ser
        self.bank = FingerSensorBank(
            sources if sources is not None else (default_source,),
            configs=configs if configs is not None else (config,), verbose=verbose,
        )
        self.max_camera_index = max_camera_index
        self.loop_video = loop_video
        self.show = show
        self.verbose = verbose
        self.read_only = read_only
        self.autostart = autostart or read_only
        self._streaming = set()
        self._should_quit = False
        self._last_search_source = None
        self._partial = bytearray()
        self._lines = deque()
        self._discard_line = False

    @property
    def multi_sensor(self):
        return len(self.bank.sources) > 1

    def reply(self, line):
        # SDK exceptions can contain newlines; every reply must be one record.
        text = " ".join(str(line).splitlines()) + "\n"
        data = text.encode("utf-8")
        if self.ser.write(data) != len(data):
            raise OSError("incomplete UART write")
        if self.verbose:
            print(f"<< {text.rstrip()}", file=sys.stderr, flush=True)

    def close_runtime(self):
        self._streaming.clear()
        try:
            self.bank.close()
        finally:
            if self.show:
                import cv2

                cv2.destroyAllWindows()

    def _poll_command_line(self, *, blocking=True):
        if self._lines:
            return self._lines.popleft()
        waiting = self.ser.in_waiting
        if not waiting and not blocking:
            return None
        # Never readline(): a timeout may return only half of a command.
        raw = self.ser.read(min(waiting, 4096) if waiting else 1)
        if self.read_only:
            return None
        for byte in raw:
            if byte in (10, 13):
                if self._discard_line:
                    self._discard_line = False
                    self._partial.clear()
                    self.reply("ERR command line too long (max 1024 bytes)")
                elif self._partial:
                    self._lines.append(self._partial.decode("utf-8", errors="replace").strip())
                    self._partial.clear()
            elif not self._discard_line:
                self._partial.append(byte)
                if len(self._partial) > 1024:
                    self._partial.clear()
                    self._discard_line = True
        return self._lines.popleft() if self._lines else None

    def _targets(self, args, *, connected=False):
        if len(args) > 1:
            raise ValueError("expected [sensor-id|all]")
        if args and args[0].lower() != "all":
            sensor_id = int(args[0])
            self.bank.validate_id(sensor_id)
            return [sensor_id]
        return sorted(self.bank.runtimes if connected else self.bank.sources)

    def _sensor_suffix(self, sensor_id):
        return f" {sensor_id}" if self.multi_sensor else ""

    def cmd_search(self, args):
        if args:
            raise ValueError("search takes no arguments")
        if self._streaming:
            raise ValueError("stop streaming before camera search")
        try:
            from orisys.capture import list_cameras

            cameras = list_cameras(max_index=self.max_camera_index)
        except Exception as exc:
            logger.exception("Camera search failed")
            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}"
            self.reply(f"camera,{cam.index},{path},{cam.width}x{cam.height}")
            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):
        # Legacy `connect <source>` targets sensor 1. `connect <id> <source>`
        # selects a slot explicitly; quote source paths containing spaces.
        if args and args[0].lower() == "all":
            if len(args) != 1:
                raise ValueError("expected connect all")
            for sensor_id in self.bank.sources:
                self._connect_one(sensor_id)
            return
        sensor_id = 1
        if len(args) >= 2 and args[0].isdigit():
            sensor_id = int(args[0])
            args = args[1:]
        source = parse_source(" ".join(args)) if args else None
        if (source is None and not self.multi_sensor and self._last_search_source is not None
                and not is_video_file(self.bank.sources[1])):
            source = self._last_search_source
        self._connect_one(sensor_id, source)

    def _connect_one(self, sensor_id, source=None):
        self.bank.validate_id(sensor_id)
        self._streaming.discard(sensor_id)
        try:
            source = self.bank.connect(sensor_id, source)
        except Exception as exc:
            logger.exception("Sensor %s connection failed", sensor_id)
            self.reply(f"ERR connect{self._sensor_suffix(sensor_id)} {exc}")
            return
        kind = "file" if is_video_file(source) else "camera"
        self.reply(f"OK connect{self._sensor_suffix(sensor_id)} {source} ({kind})")

    def cmd_run(self, args):
        targets = self._targets(args, connected=True)
        if not targets or any(i not in self.bank.runtimes for i in targets):
            raise ValueError("not connected")
        if set(targets).issubset(self._streaming):
            raise ValueError("already running")
        self._streaming.update(targets)
        self.reply("OK run" + (f" {' '.join(map(str, targets))}" if self.multi_sensor else ""))

    def cmd_stop(self, args):
        targets = self._targets(args)
        self._streaming.difference_update(targets)
        self.reply("OK stop" + (f" {' '.join(map(str, targets))}" if self.multi_sensor else ""))

    def cmd_quit(self, args):
        if args:
            raise ValueError("quit takes no arguments")
        self._streaming.clear()
        self._should_quit = True
        self.reply("OK quit")

    def _read_one(self, sensor_id, *, context):
        try:
            reading = self.bank.read(sensor_id, loop_video=self.loop_video)
        except Exception as exc:
            logger.exception("Sensor %s acquisition failed", sensor_id)
            self._streaming.discard(sensor_id)
            self.reply(f"ERR {context}{self._sensor_suffix(sensor_id)} {exc}")
            return
        if reading is None:
            self._streaming.discard(sensor_id)
            if is_video_file(self.bank.sources[sensor_id]):
                self.reply(f"OK eof{self._sensor_suffix(sensor_id)}")
            else:
                self.reply(f"ERR {context}{self._sensor_suffix(sensor_id)} grab failed")
            return
        sensor_prefix = f"{sensor_id}," if self.multi_sensor else ""
        self.reply(
            f"force,{sensor_prefix}{reading.normal:.6f},{reading.shear_x:.6f},{reading.shear_y:.6f}"
        )
        if self.show:
            self._show(reading)

    def _show(self, reading):
        import cv2
        import orisys

        tactile = self.bank.runtimes[reading.sensor_id].tactile
        try:
            arrows = orisys.util.draw_arrows(
                reading.image, tactile.flow, threshold=0.4, grid_spacing=10,
                arrow_scale=1.0, valid_mask=tactile.valid_mask,
            )
            cv2.imshow(f"UV Unrolled {reading.sensor_id}", arrows)
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                self.cmd_quit([])
            elif key == ord("r"):
                tactile.reset()
        except Exception:
            logger.exception("Sensor %s display failed", reading.sensor_id)

    def _dispatch(self, line):
        if self.read_only:
            return True
        try:
            parts = shlex.split(line)
        except ValueError as exc:
            self.reply(f"ERR command {exc}")
            return True
        if not parts:
            return True
        cmd, args = parts[0].lower(), parts[1:]
        try:
            if cmd == "search":
                self.cmd_search(args)
            elif cmd == "connect":
                self.cmd_connect(args)
            elif cmd == "run":
                self.cmd_run(args)
            elif cmd == "read":
                for sensor_id in self._targets(args):
                    self._read_one(sensor_id, context="read")
            elif cmd == "stop":
                self.cmd_stop(args)
            elif cmd in ("quit", "exit"):
                self.cmd_quit(args)
            else:
                self.reply(f"ERR unknown {cmd}")
        except (ValueError, RuntimeError) as exc:
            self.reply(f"ERR {cmd} {exc}")
        return not self._should_quit

    def serve(self):
        self.reply("OK ready")
        if self.autostart:
            self.cmd_connect(["all"])
            if self.bank.runtimes:
                self.cmd_run([])
            else:
                raise RuntimeError("autostart could not connect any sensor")
        while not self._should_quit:
            # Bounded command work keeps a noisy host from starving acquisition.
            for _ in range(16):
                line = self._poll_command_line(blocking=not self._streaming)
                if line is None:
                    break
                if not self._dispatch(line):
                    return
            for sensor_id in sorted(self._streaming):
                self._read_one(sensor_id, context="run")
                if self._should_quit:
                    return


def open_serial(port, baud, timeout, write_timeout=1.0):
    try:
        import serial
    except ImportError as exc:
        raise RuntimeError("pyserial is required: python -m pip install pyserial") from exc
    return serial.Serial(
        port=port, baudrate=baud, bytesize=serial.EIGHTBITS, parity=serial.PARITY_NONE,
        stopbits=serial.STOPBITS_ONE, timeout=timeout, write_timeout=write_timeout,
        exclusive=True,
    )


def main(argv=None):
    parser = argparse.ArgumentParser(description="UART finger control and read-only streaming")
    parser.add_argument("--port", default="/dev/ttyHS1")
    parser.add_argument("--baud", type=int, default=115200)
    parser.add_argument("--timeout", type=float, default=0.1, help="idle read timeout in seconds")
    parser.add_argument("--write-timeout", type=float, default=1.0)
    parser.add_argument("--config", "-c", action="append", help="one shared config or repeat per sensor")
    parser.add_argument("--camera-id", type=int, default=0)
    parser.add_argument("--video", help="single-sensor default video file")
    parser.add_argument("--sources", nargs="+", help="1-4 camera indices, device paths or video files")
    parser.add_argument("--max-camera-index", type=int, default=10)
    parser.add_argument("--loop-video", dest="loop_video", action="store_true", default=True)
    parser.add_argument("--no-loop-video", dest="loop_video", action="store_false")
    parser.add_argument("--show", action="store_true", help="local desktop debug only")
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--autostart", action="store_true", help="connect and run configured sources on startup")
    parser.add_argument("--read-only", action="store_true", help="autostart and ignore all UART commands")
    args = parser.parse_args(argv)
    if args.baud <= 0 or args.max_camera_index < 0:
        parser.error("baud must be positive; max camera index must be non-negative")
    if any(not math.isfinite(t) or t <= 0 for t in (args.timeout, args.write_timeout)):
        parser.error("timeouts must be finite and positive")
    if args.sources and (args.video is not None or args.camera_id != 0):
        parser.error("--sources cannot be combined with --video or --camera-id")
    configs = args.config or ["finger"]
    try:
        sources = [parse_source(s) for s in args.sources] if args.sources else [
            str(Path(args.video).resolve()) if args.video else args.camera_id,
        ]
        controller = UartFingerController(
            None, sources=sources, configs=configs, max_camera_index=args.max_camera_index,
            loop_video=args.loop_video, show=args.show, verbose=args.verbose,
            autostart=args.autostart, read_only=args.read_only,
        )
    except ValueError as exc:
        parser.error(str(exc))
    ser = None
    result = 0
    try:
        ser = open_serial(args.port, args.baud, args.timeout, args.write_timeout)
        controller.ser = ser
        print(f"UART service on {args.port} @ {args.baud}; sources={sources}", file=sys.stderr, flush=True)
        controller.serve()
    except KeyboardInterrupt:
        print("Interrupted.", file=sys.stderr)
    except Exception:
        logger.exception("UART service failed on %s", args.port)
        result = 1
    finally:
        try:
            controller.close_runtime()
        except Exception:
            logger.exception("Runtime cleanup failed")
            result = 1
        finally:
            if ser is not None:
                ser.close()
    return result


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