#!/bin/bash
# ============================================================================
#  Saqr PPE Detection - Deploy to Unitree G1
# ============================================================================
#
#  Usage (from your dev machine):
#    ./deploy.sh                    # deploy + install deps
#    ./deploy.sh --run              # deploy + install + start detection
#    ./deploy.sh --run --source 0   # deploy + start with camera 0
#
# ============================================================================

set -e

# ── Robot config ──────────────────────────────────────────────────────────────
ROBOT_IP="${ROBOT_IP:-192.168.123.164}"
ROBOT_USER="${ROBOT_USER:-unitree}"
ROBOT_ENV="${ROBOT_ENV:-teleimager}"
REMOTE_DIR="/home/${ROBOT_USER}/Saqr"
SSH_OPTS="-o StrictHostKeyChecking=no -o ConnectTimeout=10"

# ── Parse args ────────────────────────────────────────────────────────────────
RUN_AFTER=false
SOURCE="0"
MODEL="models/saqr_best.pt"
HEADLESS=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --run)       RUN_AFTER=true; shift ;;
        --source)    SOURCE="$2"; shift 2 ;;
        --model)     MODEL="$2"; shift 2 ;;
        --headless)  HEADLESS=true; shift ;;
        --ip)        ROBOT_IP="$2"; shift 2 ;;
        *)           echo "Unknown arg: $1"; exit 1 ;;
    esac
done

LOCAL_DIR="$(cd "$(dirname "$0")" && pwd)"

echo "============================================"
echo "  Saqr PPE - Deploy to Unitree G1"
echo "============================================"
echo "  Robot : ${ROBOT_USER}@${ROBOT_IP}"
echo "  Env   : ${ROBOT_ENV}"
echo "  Remote: ${REMOTE_DIR}"
echo "============================================"

# ── Step 1: Test connection ───────────────────────────────────────────────────
echo ""
echo "[1/5] Testing SSH connection..."
ssh ${SSH_OPTS} ${ROBOT_USER}@${ROBOT_IP} "echo 'Connected OK'" || {
    echo "[ERROR] Cannot reach ${ROBOT_IP}. Is the robot on?"
    exit 1
}

# ── Step 2: Create remote directory ───────────────────────────────────────────
echo "[2/5] Creating remote directory..."
ssh ${SSH_OPTS} ${ROBOT_USER}@${ROBOT_IP} "mkdir -p ${REMOTE_DIR}/{models,captures/{SAFE,PARTIAL,UNSAFE},Config,Logs}"

# ── Step 3: Copy project files ────────────────────────────────────────────────
echo "[3/5] Copying project files..."

# Python files
for f in saqr.py train.py detect.py manager.py logger.py gui.py requirements.txt; do
    if [ -f "${LOCAL_DIR}/${f}" ]; then
        scp ${SSH_OPTS} "${LOCAL_DIR}/${f}" ${ROBOT_USER}@${ROBOT_IP}:${REMOTE_DIR}/
    fi
done

# Config
scp ${SSH_OPTS} "${LOCAL_DIR}/Config/logging.json" ${ROBOT_USER}@${ROBOT_IP}:${REMOTE_DIR}/Config/

# Trained model (this is the big file)
if [ -f "${LOCAL_DIR}/models/saqr_best.pt" ]; then
    echo "  Uploading trained model (saqr_best.pt)..."
    scp ${SSH_OPTS} "${LOCAL_DIR}/models/saqr_best.pt" ${ROBOT_USER}@${ROBOT_IP}:${REMOTE_DIR}/models/
else
    echo "  [WARN] models/saqr_best.pt not found - train first!"
fi

# Base model (for retraining on robot if needed)
if [ -f "${LOCAL_DIR}/models/yolo11n.pt" ]; then
    scp ${SSH_OPTS} "${LOCAL_DIR}/models/yolo11n.pt" ${ROBOT_USER}@${ROBOT_IP}:${REMOTE_DIR}/models/
fi

echo "  Files copied."

# ── Step 4: Install dependencies ──────────────────────────────────────────────
echo "[4/5] Installing dependencies on robot..."
ssh ${SSH_OPTS} ${ROBOT_USER}@${ROBOT_IP} << 'INSTALL_EOF'
    source ~/miniconda3/etc/profile.d/conda.sh
    conda activate teleimager

    pip install -q ultralytics opencv-python numpy PyYAML 2>/dev/null
    echo "  Dependencies OK"
INSTALL_EOF

# ── Step 5: Optionally run ────────────────────────────────────────────────────
if [ "$RUN_AFTER" = true ]; then
    echo "[5/5] Starting Saqr PPE detection on robot..."

    HEADLESS_FLAG=""
    if [ "$HEADLESS" = true ]; then
        HEADLESS_FLAG="--headless"
    fi

    ssh ${SSH_OPTS} -t ${ROBOT_USER}@${ROBOT_IP} << RUN_EOF
        source ~/miniconda3/etc/profile.d/conda.sh
        conda activate teleimager
        cd ${REMOTE_DIR}
        python saqr.py --source ${SOURCE} --model ${MODEL} ${HEADLESS_FLAG}
RUN_EOF
else
    echo "[5/5] Skipped (use --run to start after deploy)"
    echo ""
    echo "============================================"
    echo "  Deployed! SSH in to run:"
    echo "============================================"
    echo ""
    echo "  ssh ${ROBOT_USER}@${ROBOT_IP}"
    echo "  conda activate teleimager"
    echo "  cd ${REMOTE_DIR}"
    echo "  python saqr.py --source 0 --model models/saqr_best.pt"
    echo ""
    echo "  Or with GUI:"
    echo "  python gui.py --source 0 --model models/saqr_best.pt"
    echo ""
fi