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

Bridge between Saqr PPE detection and the Unitree G1 arm action client.

Spawns Saqr (Project/Saqr/saqr.py) as a subprocess, parses its event stream,
and triggers the G1 'reject' arm action (id=13) whenever a tracked person
transitions to UNSAFE (= DANGER). SAFE and PARTIAL never trigger an action.

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

    # specify camera and DDS interface
    python3 saqr_g1_bridge.py --source 0 --iface enp3s0

    # dry run (no robot movement, 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
REPO_ROOT    = HERE.parent.parent                    # .../yslootahtech
SAQR_DIR     = REPO_ROOT / "Project" / "Saqr"
SAQR_SCRIPT  = SAQR_DIR / "saqr.py"

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

# 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*\|"
)


# ── G1 arm controller (lazy import: SDK only loaded when not in dry-run) ─────
class ArmController:
    def __init__(self, iface: Optional[str], timeout: float, dry_run: bool):
        self.dry_run = dry_run
        if dry_run:
            print("[BRIDGE] DRY RUN — G1 SDK will not be loaded.", flush=True)
            self.client = None
            return

        from unitree_sdk2py.core.channel import ChannelFactoryInitialize
        from unitree_sdk2py.g1.arm.g1_arm_action_client import (
            G1ArmActionClient,
            action_map,
        )
        self._action_map = action_map

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

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

    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 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.client.ExecuteAction(self._action_map[REJECT_ACTION])
        if release_after > 0:
            time.sleep(release_after)
            print(f"[BRIDGE] -> {RELEASE_ACTION}", flush=True)
            self.client.ExecuteAction(self._action_map[RELEASE_ACTION])


# ── Bridge ───────────────────────────────────────────────────────────────────
class Bridge:
    def __init__(
        self,
        arm: ArmController,
        cooldown_s: float,
        release_after_s: float,
    ):
        self.arm = arm
        self.cooldown_s = cooldown_s
        self.release_after_s = release_after_s
        self.last_status: Dict[int, str] = {}
        self.last_trigger_t: Dict[int, 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")

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

            if status != DANGER_STATUS:
                return

            # Trigger only on transitions into UNSAFE, with per-id cooldown.
            now = time.time()
            last_t = self.last_trigger_t.get(track_id, 0.0)
            transitioned = (prev != DANGER_STATUS)
            cooled_down = (now - last_t) >= self.cooldown_s

            if not (transitioned and cooled_down):
                return

            self.last_trigger_t[track_id] = now

        # Run the arm action outside the lock so we don't block parsing.
        try:
            self.arm.reject(release_after=self.release_after_s)
        except Exception as e:
            print(f"[BRIDGE][ERR] arm reject 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.")

    # 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

    arm = ArmController(iface=args.iface, timeout=args.timeout, dry_run=args.dry_run)
    bridge = Bridge(
        arm=arm,
        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()