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

H2_ROOT=${H2_ROOT:-$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")/.." && pwd)}
H2_IF=${H2_IF:-enp4s0f0}
H2_PY=${H2_PY:-python3}
H2_ADAPTER_ROOT=${H2_ADAPTER_ROOT:-$H2_ROOT/../mocap_ego_adapter}
UNITREE_ROS2_ROOT=${UNITREE_ROS2_ROOT:-$HOME/unitree_ros2}
SESSION=${H2_TMUX_SESSION:-h2-deploy}
MODE=
ACK_SUSPENDED=
ACK_JOINTS=
ACK_L1=
RELEASE_MOTION=0
TIMING_ENABLED=1
RL_MODE=off
RL_CHECKPOINT=
RL_CONFIG="$H2_ROOT/config.yaml"
EVAL_POLICY=off
EVAL_CONFIG="$H2_ROOT/configs/policy_eval.yaml"
source "$(dirname "${BASH_SOURCE[0]}")/h2_upstream_versions.sh"
h2_upstream_defaults

usage() {
  cat <<'EOF'
Usage:
  bash scripts/launch_h2_tmux.sh --dry-run
  bash scripts/launch_h2_tmux.sh --layout-smoke
  bash scripts/launch_h2_tmux.sh --live \
    --ack-suspended H2-SUSPENDED-ESTOP-READY \
    --ack-joint-ids H2-E02DC701-ID00-30-VERIFIED \
    --ack-l1-manual H2-L1-REMAINS-MANUAL

The live path creates one six-pane tmux session and stops at FROZEN_WAIT.
It never passes --start-immediately and never presses L1.
Optional: --release-motion explicitly releases high-level control before starting
the bridge; the existing suspended/E-stop acknowledgements still apply.
Optional: --no-timing disables the default diagnostic timing recorder; --timing enables it.
Optional: --deployment-version original|v1|v2 (default: original)
          --upstream-version is an alias selecting the entire adapter/predictor/policy bundle.
Optional: --rl-mode off|train|residual|baseline --rl-checkpoint /absolute/h2/checkpoint.pt
          --config /absolute/h2/config.yaml (real-ball RL remains opt-in)
Independent adapter/predictor/policy selection is not supported.
Optional: --eval-policy off|baseline --eval-config PATH
          baseline records this selected ONNX's measured-ball scores only.
          No learner, residual, exploration, time limit or automatic L1.
EOF
}

while (($#)); do
  case "$1" in
    --dry-run|--layout-smoke|--live) MODE=${1#--}; shift ;;
    --ack-suspended) ACK_SUSPENDED=${2:-}; shift 2 ;;
    --ack-joint-ids) ACK_JOINTS=${2:-}; shift 2 ;;
    --ack-l1-manual) ACK_L1=${2:-}; shift 2 ;;
    --release-motion) RELEASE_MOTION=1; shift ;;
    --timing) TIMING_ENABLED=1; shift ;;
    --no-timing) TIMING_ENABLED=0; shift ;;
    --rl-mode) RL_MODE=${2:-}; shift 2 ;;
    --rl-checkpoint) RL_CHECKPOINT=${2:-}; shift 2 ;;
    --config) RL_CONFIG=${2:-}; shift 2 ;;
    --eval-policy) EVAL_POLICY=${2:?missing evaluation mode}; shift 2 ;;
    --eval-config) EVAL_CONFIG=${2:?missing scoring config}; shift 2 ;;
    --session) SESSION=${2:-}; shift 2 ;;
    --upstream-version|--deployment-version) h2_upstream_preset "${2:-}"; shift 2 ;;
    -h|--help) usage; exit 0 ;;
    *) echo "unknown argument: $1" >&2; usage >&2; exit 2 ;;
  esac
done
h2_upstream_validate
[[ $EVAL_POLICY == off || $EVAL_POLICY == baseline ]] \
  || { echo 'invalid --eval-policy (off|baseline)' >&2; exit 2; }
if [[ $EVAL_POLICY != off ]]; then
  [[ $RL_MODE == off && -z $RL_CHECKPOINT ]] \
    || { echo 'policy evaluation cannot be combined with --rl-mode/checkpoint' >&2; exit 2; }
  [[ $H2_DEPLOYMENT_VERSION == original || $H2_DEPLOYMENT_VERSION == v2 ]] \
    || { echo 'measured-ball evaluation supports original or v2' >&2; exit 2; }
  [[ -f $EVAL_CONFIG ]] || { echo "missing scoring config: $EVAL_CONFIG" >&2; exit 2; }
fi
[[ $RL_MODE == off || $RL_MODE == train || $RL_MODE == residual || $RL_MODE == baseline ]] \
  || { echo 'invalid --rl-mode' >&2; exit 2; }
[[ $RL_MODE != residual || -n $RL_CHECKPOINT ]] \
  || { echo 'residual mode requires --rl-checkpoint' >&2; exit 2; }
[[ $RL_MODE != off || -z $RL_CHECKPOINT ]] \
  || { echo '--rl-checkpoint requires train or residual mode' >&2; exit 2; }
[[ $RL_MODE != baseline || -z $RL_CHECKPOINT ]] \
  || { echo 'baseline mode does not load a checkpoint' >&2; exit 2; }
[[ $RL_MODE == off || $H2_VELOCITY_VERSION == original ]] \
  || { echo 'RL racket reward requires --deployment-version original' >&2; exit 2; }

[[ $MODE == dry-run || $MODE == layout-smoke || $MODE == live ]] \
  || { usage >&2; exit 2; }

if [[ $MODE == dry-run ]]; then
  h2_upstream_describe
  cat <<EOF
[DRY RUN] session=$SESSION interface=$H2_IF release_motion=$RELEASE_MOTION
real-ball RL: mode=$RL_MODE config=$RL_CONFIG checkpoint=$RL_CHECKPOINT
policy evaluation: mode=$EVAL_POLICY scoring_config=$EVAL_CONFIG (actions unchanged, no training)
timing recorder: enabled=$TIMING_ENABLED (RUN_DIR/timing; no control changes)
panes: CONTROL | ADAPTER | BRIDGE | PREDICTOR | RUNNER | MONITOR
EOF
  if [[ $H2_DEPLOYMENT_VERSION == original && $RL_MODE != off ]]; then
    cat <<'EOF'
startup: simplified original RL (extended diagnostics remain manual)
gates: adapter identity -> optional explicit release -> publisher conflicts -> adapter/bridge/predictor
       -> communication ready (0.2s stable, timeout 10s) -> verified models/learner
       -> fresh CheckMode -> measured-to-frozen -> FROZEN_WAIT (runner timeout 180s)
EOF
  else
    cat <<'EOF'
gates: artifacts -> native LowState -> optional explicit release -> zero LowCmd -> ZMQ -> bridge/no-command -> planner
       -> fresh CheckMode token -> measured-to-frozen -> FROZEN_WAIT
EOF
  fi
  if [[ $EVAL_POLICY == baseline && $H2_DEPLOYMENT_VERSION == v2 ]]; then
    echo 'evaluation adapter: v2 estimator + measured racket snapshots (same transforms; record-only)'
  fi
  cat <<'EOF'
automatic L1: NEVER
active-mode release: only with --release-motion; dry-run never calls the robot
bundle options: --deployment-version original|v1|v2 (--upstream-version is an alias)
adapter/predictor/policy are always paired
EOF
  exit 0
fi

command -v tmux >/dev/null || { echo "tmux is required" >&2; exit 1; }

create_layout() {
  local placeholder='printf "pane waiting for controller...\\n"; exec sleep infinity'
  tmux new-session -d -s "$SESSION" -n deploy "bash -lc '$placeholder'"
  tmux set-option -t "$SESSION" base-index 0
  tmux set-window-option -t "$SESSION:0" pane-base-index 0
  for _ in 1 2 3 4 5; do
    tmux split-window -d -t "$SESSION:0" "bash -lc '$placeholder'"
    tmux select-layout -t "$SESSION:0" tiled >/dev/null
  done
  local titles=(CONTROL ADAPTER BRIDGE PREDICTOR RUNNER MONITOR)
  for index in 0 1 2 3 4 5; do
    tmux select-pane -t "$SESSION:0.$index" -T "${titles[$index]}"
  done
  tmux set-option -t "$SESSION" pane-border-status top
  tmux set-option -t "$SESSION" pane-border-format '#[bold] #{pane_title} #[default]'
  tmux set-option -t "$SESSION" remain-on-exit on
}

if [[ $MODE == layout-smoke ]]; then
  SESSION="h2-layout-smoke-$$"
  trap 'tmux kill-session -t "$SESSION" 2>/dev/null || true' EXIT
  create_layout
  tmux list-panes -t "$SESSION:0" -F '#{pane_index} #{pane_title}'
  [[ $(tmux list-panes -t "$SESSION:0" | wc -l) -eq 6 ]]
  echo "[PASS] six-pane tmux layout smoke test; no deployment process was started"
  exit 0
fi

[[ $ACK_SUSPENDED == H2-SUSPENDED-ESTOP-READY ]] \
  || { echo "missing exact suspended/E-stop acknowledgement" >&2; exit 1; }
[[ $ACK_JOINTS == H2-E02DC701-ID00-30-VERIFIED ]] \
  || { echo "missing exact ID 00-30 acknowledgement" >&2; exit 1; }
[[ $ACK_L1 == H2-L1-REMAINS-MANUAL ]] \
  || { echo "missing exact manual-L1 acknowledgement" >&2; exit 1; }
[[ $SESSION =~ ^[A-Za-z0-9_.-]+$ ]] || { echo "invalid tmux session name" >&2; exit 2; }
tmux has-session -t "$SESSION" 2>/dev/null \
  && { echo "tmux session already exists: $SESSION" >&2; exit 1; }

RUN_DIR="/tmp/h2_tmux_$(date +%Y%m%d_%H%M%S)_$$"
if [[ $EVAL_POLICY == baseline ]]; then
  mkdir -p "$H2_ROOT/artifacts/policy_evaluations"
  RUN_DIR="$H2_ROOT/artifacts/policy_evaluations/${RUN_DIR##*/}"
fi
mkdir -m 700 "$RUN_DIR"
printf 'H2_TIMING_ENABLED=%q\n' "$TIMING_ENABLED" > "$RUN_DIR/timing.env"
h2_upstream_check_files
h2_upstream_write_selection "$RUN_DIR"
printf 'H2_RL_MODE=%q\nH2_RL_CHECKPOINT=%q\nH2_RL_CONFIG=%q\n' \
  "$RL_MODE" "$RL_CHECKPOINT" "$RL_CONFIG" > "$RUN_DIR/rl.env"
printf '%s\n' "$RUN_DIR" > /tmp/h2_tmux_latest_run
printf 'H2_EVAL_POLICY=%q\nH2_EVAL_CONFIG=%q\nH2_EVAL_RUN_DIR=%q\n' \
  "$EVAL_POLICY" "$EVAL_CONFIG" "$RUN_DIR/evaluation" > "$RUN_DIR/eval.env"

create_layout
if [[ $EVAL_POLICY == baseline ]]; then
  printf 'H2_ROOT=%q\nH2_IF=%q\nH2_PY=%q\nH2_SESSION=%q\n' \
    "$H2_ROOT" "$H2_IF" "$H2_PY" "$SESSION" > "$RUN_DIR/eval_owner.env"
  tmux set-option -t "$SESSION" @h2_eval_run_dir "$RUN_DIR"
  roles=(control adapter bridge predictor runner monitor)
  for index in 0 1 2 3 4 5; do
    printf '%s\t%s\n' "${roles[$index]}" \
      "$(tmux display-message -p -t "$SESSION:0.$index" '#{pane_id}')" >> "$RUN_DIR/panes.tsv"
  done
  mkdir -p "$H2_ROOT/artifacts"
  printf '%s\n' "$RUN_DIR" > "$H2_ROOT/artifacts/latest_eval_run"
fi
tmux respawn-pane -k -t "$SESSION:0.0" \
  "H2_ROOT='$H2_ROOT' H2_IF='$H2_IF' H2_PY='$H2_PY' H2_ADAPTER_ROOT='$H2_ADAPTER_ROOT' UNITREE_ROS2_ROOT='$UNITREE_ROS2_ROOT' bash '$H2_ROOT/scripts/h2_tmux_orchestrator.sh' '$SESSION' '$RUN_DIR' '$RELEASE_MOTION'"
tmux select-pane -t "$SESSION:0.0"

echo "Attaching tmux session $SESSION; logs: $RUN_DIR"
if [[ -n ${TMUX:-} ]]; then
  tmux switch-client -t "$SESSION"
  exit 0
fi
exec tmux attach-session -t "$SESSION"