#!/usr/bin/env python3
"""
saqr_g1_bridge.py

Bridge between Saqr PPE detection and the Unitree G1 robot.

Spawns Saqr (saqr.py in this same folder) as a subprocess, parses its event
stream, and on each per-person status transition:

  * UNSAFE  -> announce "Not safe!" via the G1 onboard TtsMaker (English,
              speaker_id=2) AND run the 'reject' arm action (id=13).
  * SAFE    -> announce "Safe!" via the G1 onboard TtsMaker. No arm motion.
  * PARTIAL -> nothing.

Both DDS clients (G1ArmActionClient and G1 AudioClient) share a single
ChannelFactoryInitialize call. The TTS speaker_id was identified by running
Project/Sanad/voice_example.py mode 6 — speaker_id=2 is English on current
G1 firmware (speaker_id=0 is Chinese regardless of input text).

Saqr event line format (from emit_event in saqr.py):
    ID 0001 | NEW | UNSAFE | wearing: ... | missing: ... | unknown: ...
    ID 0001 | STATUS_CHANGE | SAFE  | wearing: ... | missing: ... | unknown: ...

Usage:
    # default: webcam, default DDS interface
    python3 saqr_g1_bridge.py

    # on the robot
    python3 saqr_g1_bridge.py --iface eth0 --source realsense --headless

    # dry run (no robot movement / TTS, just print decisions)
    python3 saqr_g1_bridge.py --dry-run

    # forward extra args to saqr.py after a `--`
    python3 saqr_g1_bridge.py --iface eth0 -- --conf 0.4 --imgsz 640
"""

from __future__ import annotations

import argparse
import os
import re
import signal
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Dict, Optional


# ── Defaults ─────────────────────────────────────────────────────────────────
HERE         = Path(__file__).resolve().parent
SAQR_DIR     = HERE                                  # bridge lives next to saqr.py
SAQR_SCRIPT  = SAQR_DIR / "saqr.py"

DANGER_STATUS = "UNSAFE"
SAFE_STATUS   = "SAFE"
REJECT_ACTION = "reject"
RELEASE_ACTION = "release arm"

# G1 onboard TtsMaker (see Project/Sanad/voice_example.py mode 6).
# speaker_id=2 was confirmed English on current G1 firmware.
TTS_SPEAKER_ID = 2
TTS_VOLUME     = 100

TTS_TEXT_SAFE            = "Safe to enter. Have a good day."
TTS_UNSAFE_WITH_MISSING  = (
    "Please stop. Wear your proper safety equipment. You are missing {items}."
)
TTS_UNSAFE_GENERIC       = (
    "Please stop. Wear your proper safety equipment."
)

# ID NNNN | EVENT_TYPE | STATUS | wearing: ... | missing: ... | unknown: ...
EVENT_RE = re.compile(
    r"^ID\s+(?P<id>\d+)\s*\|\s*"
    r"(?P<event>NEW|STATUS_CHANGE)\s*\|\s*"
    r"(?P<status>SAFE|PARTIAL|UNSAFE)\s*\|\s*"
    r"wearing:\s*(?P<wearing>[^|]*?)\s*\|\s*"
    r"missing:\s*(?P<missing>[^|]*?)\s*\|\s*"
    r"unknown:\s*(?P<unknown>.*?)\s*$"
)


def _parse_list_field(s: str) -> list:
    """Parse 'helmet, vest' or 'none' into a list of items."""
    s = (s or "").strip()
    if not s or s.lower() == "none":
        return []
    return [x.strip() for x in s.split(",") if x.strip()]


def _human_join(items: list) -> str:
    """Join a list in natural English: 'helmet and vest', 'a, b, and c'."""
    if not items:
        return ""
    if len(items) == 1:
        return items[0]
    if len(items) == 2:
        return f"{items[0]} and {items[1]}"
    return ", ".join(items[:-1]) + f", and {items[-1]}"


def build_unsafe_tts(missing: list) -> str:
    if not missing:
        return TTS_UNSAFE_GENERIC
    return TTS_UNSAFE_WITH_MISSING.format(items=_human_join(missing))


# ── G1 robot controller (lazy import: SDK only loaded when not in dry-run) ───
class RobotController:
    """Owns both the G1 arm action client and the G1 audio (TTS) client.

    A single ChannelFactoryInitialize call is shared by both clients.
    """

    def __init__(
        self,
        iface: Optional[str],
        timeout: float,
        dry_run: bool,
        tts_speaker_id: int,
    ):
        self.dry_run = dry_run
        self.tts_speaker_id = tts_speaker_id
        self.arm_client = None
        self.audio_client = None
        self._action_map = None

        if dry_run:
            print("[BRIDGE] DRY RUN — G1 SDK will not be loaded.", flush=True)
            return

        from unitree_sdk2py.core.channel import ChannelFactoryInitialize
        from unitree_sdk2py.g1.arm.g1_arm_action_client import (
            G1ArmActionClient,
            action_map,
        )
        from unitree_sdk2py.g1.audio.g1_audio_client import AudioClient

        self._action_map = action_map

        if iface:
            ChannelFactoryInitialize(0, iface)
        else:
            ChannelFactoryInitialize(0)

        self.arm_client = G1ArmActionClient()
        self.arm_client.SetTimeout(timeout)
        self.arm_client.Init()
        print(f"[BRIDGE] G1ArmActionClient ready (iface={iface or 'default'})",
              flush=True)

        self.audio_client = AudioClient()
        self.audio_client.SetTimeout(timeout)
        self.audio_client.Init()
        try:
            self.audio_client.SetVolume(TTS_VOLUME)
        except Exception as e:
            print(f"[BRIDGE][WARN] AudioClient.SetVolume failed: {e}", flush=True)
        print(f"[BRIDGE] G1 AudioClient ready (speaker_id={tts_speaker_id})",
              flush=True)

    # ── TTS ─────────────────────────────────────────────────────────────────
    def speak(self, text: str):
        if self.dry_run:
            print(f"[BRIDGE] (dry) would TtsMaker({text!r}, "
                  f"speaker_id={self.tts_speaker_id})", flush=True)
            return
        if self.audio_client is None:
            return
        try:
            print(f"[BRIDGE] tts -> {text!r}", flush=True)
            code = self.audio_client.TtsMaker(text, self.tts_speaker_id)
            if code != 0:
                print(f"[BRIDGE][WARN] TtsMaker return code = {code}", flush=True)
        except Exception as e:
            print(f"[BRIDGE][ERR] TtsMaker failed: {e}", flush=True)

    # ── Arm ─────────────────────────────────────────────────────────────────
    def reject(self, release_after: float):
        if self.dry_run:
            print(f"[BRIDGE] (dry) would run '{REJECT_ACTION}' "
                  f"then release after {release_after:.1f}s", flush=True)
            return
        if self.arm_client is None or self._action_map is None:
            return
        if REJECT_ACTION not in self._action_map:
            print(f"[BRIDGE][ERR] '{REJECT_ACTION}' not in SDK action_map",
                  flush=True)
            return
        print(f"[BRIDGE] -> {REJECT_ACTION}", flush=True)
        self.arm_client.ExecuteAction(self._action_map[REJECT_ACTION])
        if release_after > 0:
            time.sleep(release_after)
            print(f"[BRIDGE] -> {RELEASE_ACTION}", flush=True)
            self.arm_client.ExecuteAction(self._action_map[RELEASE_ACTION])


# ── Bridge ───────────────────────────────────────────────────────────────────
class Bridge:
    def __init__(
        self,
        robot: RobotController,
        cooldown_s: float,
        release_after_s: float,
    ):
        self.robot = robot
        self.cooldown_s = cooldown_s
        self.release_after_s = release_after_s
        self.last_status: Dict[int, str] = {}
        # Per-id cooldown is keyed by (track_id, status) so a SAFE announce
        # and an UNSAFE announce don't share the same timer.
        self.last_trigger_t: Dict[tuple[int, str], float] = {}
        self._lock = threading.Lock()

    def handle_line(self, line: str):
        line = line.rstrip()
        if not line:
            return
        # Always echo Saqr output so the user still sees the live stream.
        print(line, flush=True)

        m = EVENT_RE.match(line)
        if not m:
            return

        track_id = int(m.group("id"))
        status = m.group("status")
        missing = _parse_list_field(m.group("missing"))

        with self._lock:
            prev = self.last_status.get(track_id)
            self.last_status[track_id] = status

            # Only SAFE / UNSAFE transitions trigger the robot. PARTIAL is silent.
            if status not in (DANGER_STATUS, SAFE_STATUS):
                return

            # Only fire on transitions, not on every NEW/STATUS_CHANGE for the
            # same status.
            if prev == status:
                return

            now = time.time()
            last_t = self.last_trigger_t.get((track_id, status), 0.0)
            if (now - last_t) < self.cooldown_s:
                return
            self.last_trigger_t[(track_id, status)] = now

        # Run robot actions outside the lock so we don't block parsing.
        try:
            if status == DANGER_STATUS:
                self.robot.speak(build_unsafe_tts(missing))
                self.robot.reject(release_after=self.release_after_s)
            else:  # SAFE
                self.robot.speak(TTS_TEXT_SAFE)
        except Exception as e:
            print(f"[BRIDGE][ERR] robot action failed: {e}", flush=True)


# ── Saqr subprocess management ───────────────────────────────────────────────
def build_saqr_cmd(saqr_extra_args: list[str]) -> list[str]:
    if not SAQR_SCRIPT.exists():
        sys.exit(f"[BRIDGE][FATAL] saqr.py not found at: {SAQR_SCRIPT}")
    # -u for unbuffered stdout (so events arrive line-by-line).
    return [sys.executable, "-u", str(SAQR_SCRIPT), *saqr_extra_args]


def split_argv(argv: list[str]) -> tuple[list[str], list[str]]:
    """Split bridge args from saqr passthrough args at the first '--'."""
    if "--" in argv:
        idx = argv.index("--")
        return argv[:idx], argv[idx + 1 :]
    return argv, []


def main():
    bridge_argv, saqr_extra = split_argv(sys.argv[1:])

    ap = argparse.ArgumentParser(
        description="Bridge Saqr PPE events to the G1 arm 'reject' action."
    )
    ap.add_argument("--iface", default=None,
                    help="DDS network interface (e.g. enp3s0). Optional.")
    ap.add_argument("--timeout", type=float, default=10.0,
                    help="G1 arm client timeout (seconds).")
    ap.add_argument("--cooldown", type=float, default=8.0,
                    help="Per-track-id seconds before reject can re-trigger.")
    ap.add_argument("--release-after", type=float, default=2.0,
                    help="Seconds before auto-running 'release arm' (0 = never).")
    ap.add_argument("--dry-run", action="store_true",
                    help="Parse and decide but never call the SDK.")
    ap.add_argument("--speaker-id", type=int, default=TTS_SPEAKER_ID,
                    help=f"G1 TtsMaker speaker_id (default {TTS_SPEAKER_ID}, English).")

    # Convenience pass-throughs to saqr.py (you can also use `-- ...`).
    ap.add_argument("--source", default=None,
                    help="Saqr --source (0/realsense/path). Default: leave to saqr.")
    ap.add_argument("--headless", action="store_true",
                    help="Pass --headless to saqr.")
    ap.add_argument("--saqr-conf", type=float, default=None,
                    help="Pass --conf to saqr.")
    ap.add_argument("--imgsz", type=int, default=None,
                    help="Pass --imgsz to saqr.")
    ap.add_argument("--device", default=None,
                    help="Pass --device to saqr (e.g. cpu / 0 / cuda:0).")

    args = ap.parse_args(bridge_argv)

    # Build saqr args from convenience flags + raw passthrough.
    saqr_args: list[str] = []
    if args.source is not None:
        saqr_args += ["--source", args.source]
    if args.headless:
        saqr_args += ["--headless"]
    if args.saqr_conf is not None:
        saqr_args += ["--conf", str(args.saqr_conf)]
    if args.imgsz is not None:
        saqr_args += ["--imgsz", str(args.imgsz)]
    if args.device is not None:
        saqr_args += ["--device", args.device]
    saqr_args += saqr_extra

    robot = RobotController(
        iface=args.iface,
        timeout=args.timeout,
        dry_run=args.dry_run,
        tts_speaker_id=args.speaker_id,
    )
    bridge = Bridge(
        robot=robot,
        cooldown_s=args.cooldown,
        release_after_s=args.release_after,
    )

    cmd = build_saqr_cmd(saqr_args)
    print(f"[BRIDGE] launching: {' '.join(cmd)}", flush=True)
    print(f"[BRIDGE] cwd: {SAQR_DIR}", flush=True)

    env = os.environ.copy()
    env["PYTHONUNBUFFERED"] = "1"

    proc = subprocess.Popen(
        cmd,
        cwd=str(SAQR_DIR),
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        bufsize=1,
        text=True,
        env=env,
    )

    def _forward_signal(signum, _frame):
        print(f"[BRIDGE] signal {signum} -> stopping saqr", flush=True)
        try:
            proc.send_signal(signum)
        except Exception:
            pass

    signal.signal(signal.SIGINT, _forward_signal)
    signal.signal(signal.SIGTERM, _forward_signal)

    try:
        assert proc.stdout is not None
        for line in proc.stdout:
            bridge.handle_line(line)
    finally:
        rc = proc.wait()
        print(f"[BRIDGE] saqr exited rc={rc}", flush=True)
        sys.exit(rc)


if __name__ == "__main__":
    main()