#!/usr/bin/env python3
"""Headless UART finger service for 1-4 sensors.

The login console and data UART must be separate ports. Cameras are selected
explicitly by UART commands or a startup setup file. All force records include
a sensor ID. Importing this module opens no devices; acquisition is owned by
finger_api.

Commands (one line, case-insensitive, CR or LF terminated):
  add SOURCE [CONFIG] [--calibration PATH]
  load SETUP.json
  remove SENSOR_ID
  start | stop | read | quit | exit
  list | status | help

Replies:
  OK ready <n>
  OK added <id> <kind> <source>
  OK removed <id>
  OK loaded <id> [<id> ...]
  OK start | OK stop | OK quit
  OK list <n>
  sensor,<id>,<kind>,<source>,<config>,<calibration>
  OK status sensors=<n> streaming=<0|1> loop_video=<0|1>
  force,<id>,<fnormal>,<fshearx>,<fsheary>
  OK eof <context>
  ERR <context> <message>
"""

from __future__ import annotations

import argparse
import difflib
import logging
import math
import os
import shlex
import sys
import time
from collections import deque

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

from finger_api import FingerSensorBank, is_video_file, parse_source

logger = logging.getLogger(__name__)

_KNOWN_COMMANDS = (
    "add", "load", "remove", "start", "stop", "read", "quit", "exit",
    "list", "status", "help",
)

_HELP_LINES = (
    "add SOURCE [CONFIG] [--calibration PATH]",
    "load SETUP.json",
    "remove SENSOR_ID",
    "start | stop | read | quit",
    "list | status | help",
)


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

    def __init__(
        self, ser, *, config="finger", loop_video=True,
        show=False, verbose=False, max_hz=0.0, prompt=False,
    ):
        self.ser = ser
        self.bank = FingerSensorBank(config=config, verbose=verbose)
        self.loop_video = loop_video
        self.show = show
        self.verbose = verbose
        self.max_hz = max_hz
        self.prompt = prompt
        self._streaming = False
        self._should_quit = False
        self._partial = bytearray()
        self._lines = deque()
        self._discard_line = False

    # ------------------------------------------------------------------ io

    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")
        self.ser.flush()
        if self.verbose:
            print(f"<< {text.rstrip()}", file=sys.stderr, flush=True)

    def _maybe_prompt(self):
        if self.prompt and not self._streaming:
            self.reply(">")

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

                cv2.destroyAllWindows()

    def _poll_command_line(self, *, blocking=True):
        if self._lines:
            line = self._lines.popleft()
            if self.verbose:
                print(f">> {line}", file=sys.stderr, flush=True)
            return line
        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)
        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
        if self._lines:
            line = self._lines.popleft()
            if self.verbose:
                print(f">> {line}", file=sys.stderr, flush=True)
            return line
        return None

    # ------------------------------------------------------------- streaming

    def _read_round(self):
        try:
            readings = self.bank.read_all(loop_video=self.loop_video)
        except EOFError as exc:
            self._streaming = False
            self.reply(f"OK eof {exc}")
            self._maybe_prompt()
            return
        except Exception as exc:
            logger.exception("Acquisition round failed")
            self._streaming = False
            self.reply(f"ERR read {exc}")
            self._maybe_prompt()
            return
        for reading in readings:
            self.reply(
                f"force,{reading.sensor_id},{reading.normal:.6f},"
                f"{reading.shear_x:.6f},{reading.shear_y:.6f}"
            )
            if self.show:
                self._show(reading)
                if self._should_quit:
                    return

    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._streaming = False
                self._should_quit = True
                self.reply("OK quit")
            elif key == ord("r"):
                tactile.reset()
        except Exception:
            logger.exception("Sensor %s display failed", reading.sensor_id)

    # ------------------------------------------------------------- dispatch

    def _reply_hint(self, command: str):
        close = difflib.get_close_matches(command, _KNOWN_COMMANDS, n=1)
        suffix = f" (did you mean '{close[0]}'?)" if close else ""
        self.reply(f"ERR unknown {command}{suffix}")

    def _cmd_add(self, parts):
        if self._streaming:
            self.reply("ERR add stop streaming first")
            return
        parser = argparse.ArgumentParser(prog="add", add_help=False, exit_on_error=False)
        parser.add_argument("source")
        parser.add_argument("config", nargs="?")
        parser.add_argument("--calibration")
        try:
            args = parser.parse_args(parts)
        except (argparse.ArgumentError, SystemExit):
            self.reply("ERR command expected add SOURCE [CONFIG] [--calibration PATH]")
            return
        try:
            sensor_id = self.bank.add(
                parse_source(args.source), config=args.config, cal_path=args.calibration,
            )
        except FileNotFoundError:
            self.reply(f"ERR add source not found: {args.source}")
        except ValueError as exc:
            self.reply(f"ERR add {exc}")
        except Exception as exc:
            logger.exception("Sensor addition failed")
            self.reply(
                f"ERR add failed for {args.source}: {type(exc).__name__}: {exc}"
            )
        else:
            src = self.bank.sources[sensor_id]
            kind = "file" if is_video_file(src) else "camera"
            self.reply(f"OK added {sensor_id} {kind} {src}")

    def _cmd_load(self, parts):
        if len(parts) != 1:
            self.reply("ERR command expected load SETUP.json")
            return
        if self._streaming:
            self.reply("ERR load stop streaming first")
            return
        try:
            sensor_ids = self.bank.add_setup(parts[0])
        except FileNotFoundError:
            self.reply(f"ERR load setup not found: {parts[0]}")
        except ValueError as exc:
            self.reply(f"ERR load {exc}")
        except Exception as exc:
            logger.exception("Setup loading failed")
            self.reply(f"ERR load failed: {type(exc).__name__}: {exc}")
        else:
            self.reply("OK loaded " + " ".join(map(str, sensor_ids)))

    def _cmd_remove(self, parts):
        if self._streaming:
            self.reply("ERR remove stop streaming first")
            return
        if len(parts) != 1 or not parts[0].isdigit():
            self.reply("ERR command expected remove SENSOR_ID")
            return
        sensor_id = int(parts[0])
        try:
            self.bank.remove_sensor(sensor_id)
        except ValueError as exc:
            self.reply(f"ERR remove {exc}")
        except Exception as exc:
            logger.exception("Sensor removal failed")
            self.reply(f"ERR remove failed: {type(exc).__name__}: {exc}")
        else:
            self.reply(f"OK removed {sensor_id}")

    def _cmd_list(self):
        for info in self.bank.list_sensors():
            cal = info.calibration or "-"
            self.reply(
                f"sensor,{info.sensor_id},{info.kind},{info.source},{info.config},{cal}"
            )
        self.reply(f"OK list {len(self.bank.sources)}")

    def _cmd_status(self):
        self.reply(
            f"OK status sensors={len(self.bank.sources)} "
            f"streaming={int(self._streaming)} "
            f"loop_video={int(self.loop_video)}"
        )

    def _cmd_help(self):
        self.reply("OK help " + " | ".join(_HELP_LINES))

    def _cmd_start(self, parts):
        if parts:
            self.reply("ERR command start takes no arguments")
            return
        if not self.bank.runtimes:
            self.reply("ERR start no sensors added")
            return
        self._streaming = True
        self.reply("OK start")

    def _cmd_stop(self, parts):
        if parts:
            self.reply("ERR command stop takes no arguments")
            return
        self._streaming = False
        self.reply("OK stop")
        self._maybe_prompt()

    def _cmd_read(self, parts):
        if parts:
            self.reply("ERR command read takes no arguments")
            return
        if not self.bank.runtimes:
            self.reply("ERR read no sensors added")
            return
        self._read_round()

    def _cmd_quit(self, parts):
        if parts:
            self.reply("ERR command quit takes no arguments")
            return
        self._streaming = False
        self._should_quit = True
        self.reply("OK quit")

    def _dispatch(self, line):
        try:
            parts = shlex.split(line)
        except ValueError as exc:
            self.reply(f"ERR command {exc}")
            self._maybe_prompt()
            return True
        if not parts:
            self._maybe_prompt()
            return True
        command = parts[0].lower()
        args = parts[1:]

        if command == "add":
            self._cmd_add(args)
        elif command == "load":
            self._cmd_load(args)
        elif command == "remove":
            self._cmd_remove(args)
        elif command == "start":
            self._cmd_start(args)
        elif command == "stop":
            self._cmd_stop(args)
        elif command == "read":
            self._cmd_read(args)
        elif command in ("quit", "exit"):
            self._cmd_quit(args)
        elif command == "list":
            self._cmd_list()
        elif command == "status":
            self._cmd_status()
        elif command in ("help", "?"):
            self._cmd_help()
        else:
            self._reply_hint(command)

        self._maybe_prompt()
        return not self._should_quit

    # ---------------------------------------------------------------- serve

    def serve(self, *, start=False):
        if start and not self.bank.runtimes:
            raise RuntimeError("no sensors added")
        self.reply(f"OK ready {len(self.bank.sources)}")
        self._maybe_prompt()
        self._streaming = start
        while not self._should_quit:
            for _ in range(16):
                line = self._poll_command_line(blocking=not self._streaming)
                if line is None:
                    break
                if not self._dispatch(line):
                    return
            if self._streaming:
                t0 = time.monotonic()
                self._read_round()
                if self.max_hz > 0 and self._streaming:
                    budget = 1.0 / self.max_hz - (time.monotonic() - t0)
                    if budget > 0:
                        time.sleep(budget)


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="Manual UART finger service: add, load, remove, start, stop, read, list, status, quit"
    )
    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", default="finger", help="default config for add commands")
    parser.add_argument("--setup", help="JSON file listing sources, configs and calibrations to connect")
    parser.add_argument("--start", action="store_true", help="start output after loading --setup")
    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(
        "--max-hz", type=float, default=0.0,
        help="cap force rounds per second while streaming (0 = unlimited)",
    )
    parser.add_argument(
        "--prompt", action="store_true",
        help="emit a '>' line before each idle command wait (off by default)",
    )
    args = parser.parse_args(argv)
    if args.start and not args.setup:
        parser.error("--start requires --setup")
    if args.baud <= 0:
        parser.error("baud must be positive")
    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 not math.isfinite(args.max_hz) or args.max_hz < 0:
        parser.error("--max-hz must be finite and non-negative")
    controller = UartFingerController(
        None, config=args.config, loop_video=args.loop_video,
        show=args.show, verbose=args.verbose, max_hz=args.max_hz, prompt=args.prompt,
    )
    ser = None
    result = 0
    try:
        if args.setup:
            controller.bank.add_setup(args.setup)
        ser = open_serial(args.port, args.baud, args.timeout, args.write_timeout)
        controller.ser = ser
        print(
            f"UART service on {args.port} @ {args.baud}; "
            f"{len(controller.bank.runtimes)} sensors ready",
            file=sys.stderr, flush=True,
        )
        controller.serve(start=args.start)
    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())