#!/usr/bin/env bash
set -euo pipefail

ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "$ROOT"

TASK="${TASK:-}"
if [[ -n "$TASK" ]]; then
  TASK_ENV="$ROOT/deployment/tasks/${TASK}.env"
  LOCAL_TASK_ENV="$ROOT/local_pi_mamba/deployment_tasks/${TASK}.env"
  if [[ ! -f "$TASK_ENV" ]]; then
    if [[ -f "$LOCAL_TASK_ENV" ]]; then
      TASK_ENV="$LOCAL_TASK_ENV"
    else
      echo "Unknown task preset: $TASK" >&2
      echo "Expected: $TASK_ENV" >&2
      echo "      or: $LOCAL_TASK_ENV" >&2
      exit 2
    fi
  fi
  source "$TASK_ENV"
fi

PYTHON="${PYTHON:-python}"
OPENPI_REPO="${OPENPI_REPO:-${HOME}/projects/openpi}"
TELEOP_REPO="${TELEOP_REPO:-${ROOT}/UR_SpaceMouse_Teleop}"
ROBOT_IP="${ROBOT_IP:-192.168.56.101}"
HOST="${HOST:-127.0.0.1}"
PORT="${PORT:-8018}"
PROMPT="${PROMPT:-}"
ACTION_MODE="${ACTION_MODE:-joint}"
GRIPPER_MODE="${GRIPPER_MODE:-onrobot-rg2}"
SET_TCP="${SET_TCP:-1}"
TCP_Z_MM="${TCP_Z_MM:-240}"
OLD_TCP_Z_MM="${OLD_TCP_Z_MM:-199}"
TCP_COMPENSATION_SIGN="${TCP_COMPENSATION_SIGN:--1}"
CAMERA_SERIALS="${CAMERA_SERIALS:-405622073020,346122070666}"
CAMERA_WIDTH="${CAMERA_WIDTH:-640}"
CAMERA_HEIGHT="${CAMERA_HEIGHT:-480}"
CAMERA_FPS="${CAMERA_FPS:-30}"
CAMERA_STARTUP_TIMEOUT="${CAMERA_STARTUP_TIMEOUT:-20}"
PRIMARY_CAMERA_INDEX="${PRIMARY_CAMERA_INDEX:-0}"
WRIST_CAMERA_INDEX="${WRIST_CAMERA_INDEX:-1}"
NO_DEPTH="${NO_DEPTH:-1}"
CONTROL_FREQUENCY="${CONTROL_FREQUENCY:-20}"
REPLAN_STEPS="${REPLAN_STEPS:-30}"
RTC="${RTC:-0}"
RTC_EXECUTION_HORIZON="${RTC_EXECUTION_HORIZON:-${RTC_STEPS:-0}}"
RTC_MAX_GUIDANCE_WEIGHT="${RTC_MAX_GUIDANCE_WEIGHT:-${RTC_GUIDANCE_WEIGHT:-5.0}}"
RTC_PREFIX_ATTENTION_SCHEDULE="${RTC_PREFIX_ATTENTION_SCHEDULE:-${RTC_SCHEDULE:-exp}}"
RTC_QUERY_AT_STEPS="${RTC_QUERY_AT_STEPS:-0}"
RTC_DELAY_BUFFER_SIZE="${RTC_DELAY_BUFFER_SIZE:-10}"
RTC_LITE="${RTC_LITE:-0}"
RTC_LITE_STEPS="${RTC_LITE_STEPS:-10}"
RTC_LITE_SCHEDULE="${RTC_LITE_SCHEDULE:-exp}"
MAX_JOINT_STEP_RAD="${MAX_JOINT_STEP_RAD:-0.6}"
MAX_POS_STEP_M="${MAX_POS_STEP_M:-0.003}"
MAX_ROT_STEP_RAD="${MAX_ROT_STEP_RAD:-0.03}"
SERVO_VELOCITY="${SERVO_VELOCITY:-0.10}"
SERVO_ACCELERATION="${SERVO_ACCELERATION:-0.20}"
SERVO_LOOKAHEAD_TIME="${SERVO_LOOKAHEAD_TIME:-0.1}"
SERVO_GAIN="${SERVO_GAIN:-100}"
INTERACTIVE_TRIALS="${INTERACTIVE_TRIALS:-0}"
START_PAUSED="${START_PAUSED:-0}"
START_POSE_NAME="${START_POSE_NAME:-}"
START_POSES_ROOT="${START_POSES_ROOT:-${ROOT}/robot_poses}"
RESET_SPEED="${RESET_SPEED:-0.45}"
RESET_ACCELERATION="${RESET_ACCELERATION:-0.45}"
OPEN_GRIPPER_ON_RESET="${OPEN_GRIPPER_ON_RESET:-0}"
RECORD_RGB_VIDEO="${RECORD_RGB_VIDEO:-0}"
RECORD_TASK_NAME="${RECORD_TASK_NAME:-${TASK_NAME:-${TASK:-task}}}"
RGB_VIDEO_ROOT="${RGB_VIDEO_ROOT:-}"
RGB_VIDEO_CAMERA_NAMES="${RGB_VIDEO_CAMERA_NAMES:-front,wrist,top}"
RGB_VIDEO_CODEC="${RGB_VIDEO_CODEC:-mp4v}"
RGB_VIDEO_FPS="${RGB_VIDEO_FPS:-0}"
RGB_VIDEO_TRANSCODE_H264="${RGB_VIDEO_TRANSCODE_H264:-1}"
EXECUTE="${EXECUTE:-0}"
REPORT_ROOT="${REPORT_ROOT:-${ROOT}/reports}"

DURATION_VALUE="${DURATION:-60}"
if [[ "${INTERACTIVE_TRIALS}" == 1 && -z "${DURATION+x}" ]]; then
  DURATION_VALUE=""
fi

if [[ -z "$PROMPT" ]]; then
  echo "Set PROMPT=<task prompt>, or set TASK=<preset>." >&2
  exit 2
fi
if ! PYTHON_RESOLVED="$(command -v "$PYTHON" 2>/dev/null)"; then
  echo "Python not found: $PYTHON" >&2
  exit 2
fi
PYTHON="$PYTHON_RESOLVED"

"$PYTHON" - <<'PY'
import importlib
import sys

missing = []
for module in ("websockets", "rtde_control", "pyrealsense2", "scipy", "numpy"):
    try:
        importlib.import_module(module)
    except Exception as exc:
        missing.append(f"{module}: {exc}")

if missing:
    print(f"Robot deployment Python is missing required packages: {sys.executable}", file=sys.stderr)
    for item in missing:
        print(f"  - {item}", file=sys.stderr)
    raise SystemExit(2)
PY

args=(
  "$ROOT/deployment/scripts/real_robot_inference.py"
  --host "$HOST"
  --port "$PORT"
  --robot-ip "$ROBOT_IP"
  --prompt "$PROMPT"
  --teleop-repo "$TELEOP_REPO"
  --openpi-repo "$OPENPI_REPO"
  --output-dir "$REPORT_ROOT"
  --action-mode "$ACTION_MODE"
  --control-frequency "$CONTROL_FREQUENCY"
  --replan-steps "$REPLAN_STEPS"
  --rtc-execution-horizon "$RTC_EXECUTION_HORIZON"
  --rtc-max-guidance-weight "$RTC_MAX_GUIDANCE_WEIGHT"
  --rtc-prefix-attention-schedule "$RTC_PREFIX_ATTENTION_SCHEDULE"
  --rtc-query-at-steps "$RTC_QUERY_AT_STEPS"
  --rtc-delay-buffer-size "$RTC_DELAY_BUFFER_SIZE"
  --rtc-lite-steps "$RTC_LITE_STEPS"
  --rtc-lite-schedule "$RTC_LITE_SCHEDULE"
  --gripper-mode "$GRIPPER_MODE"
  --tcp-z-mm "$TCP_Z_MM"
  --old-tcp-z-mm "$OLD_TCP_Z_MM"
  --tcp-compensation-sign "$TCP_COMPENSATION_SIGN"
  --max-pos-step-m "$MAX_POS_STEP_M"
  --max-rot-step-rad "$MAX_ROT_STEP_RAD"
  --max-joint-step-rad "$MAX_JOINT_STEP_RAD"
  --servo-velocity "$SERVO_VELOCITY"
  --servo-acceleration "$SERVO_ACCELERATION"
  --servo-lookahead-time "$SERVO_LOOKAHEAD_TIME"
  --servo-gain "$SERVO_GAIN"
  --start-poses-root "$START_POSES_ROOT"
  --reset-speed "$RESET_SPEED"
  --reset-acceleration "$RESET_ACCELERATION"
  --camera-serial "$CAMERA_SERIALS"
  --camera-width "$CAMERA_WIDTH"
  --camera-height "$CAMERA_HEIGHT"
  --camera-fps "$CAMERA_FPS"
  --camera-startup-timeout "$CAMERA_STARTUP_TIMEOUT"
  --primary-camera-index "$PRIMARY_CAMERA_INDEX"
  --wrist-camera-index "$WRIST_CAMERA_INDEX"
  --record-task-name "$RECORD_TASK_NAME"
  --rgb-video-camera-names "$RGB_VIDEO_CAMERA_NAMES"
  --rgb-video-codec "$RGB_VIDEO_CODEC"
  --rgb-video-fps "$RGB_VIDEO_FPS"
)

if [[ -n "$DURATION_VALUE" ]]; then
  args+=(--duration "$DURATION_VALUE")
fi
if [[ "$SET_TCP" == 1 ]]; then
  args+=(--set-tcp)
fi
if [[ "$RTC" == 1 ]]; then
  args+=(--rtc)
fi
if [[ "$RTC_LITE" == 1 ]]; then
  args+=(--rtc-lite)
fi
if [[ "$INTERACTIVE_TRIALS" == 1 ]]; then
  args+=(--interactive-trials)
fi
if [[ "$START_PAUSED" == 1 ]]; then
  args+=(--start-paused)
fi
if [[ -n "$START_POSE_NAME" ]]; then
  args+=(--start-pose-name "$START_POSE_NAME")
fi
if [[ "$OPEN_GRIPPER_ON_RESET" == 1 ]]; then
  args+=(--open-gripper-on-reset)
fi
if [[ "$RECORD_RGB_VIDEO" == 1 ]]; then
  args+=(--record-rgb-videos)
fi
if [[ "$RGB_VIDEO_TRANSCODE_H264" == 1 ]]; then
  args+=(--rgb-video-transcode-h264)
fi
if [[ -n "$RGB_VIDEO_ROOT" ]]; then
  args+=(--rgb-video-root "$RGB_VIDEO_ROOT")
fi
if [[ "$NO_DEPTH" == 1 ]]; then
  args+=(--no-depth)
fi
if [[ "$EXECUTE" == 1 ]]; then
  args+=(--execute)
fi

exec "$PYTHON" "${args[@]}"