#!/bin/bash
# ============================================================================
#  Saqr PPE Detection - Deploy to Unitree G1
# ============================================================================
#
#  Usage (from dev machine, any cwd):
#    scripts/deploy.sh                    # deploy + install deps
#    scripts/deploy.sh --run              # deploy + install + start bridge
#    scripts/deploy.sh --ip 10.0.0.5      # custom robot IP
#
# ============================================================================

set -e

HERE="$(cd "$(dirname "$0")" && pwd)"
PROJECT_ROOT="$(cd "$HERE/.." && pwd)"

ROBOT_IP="${ROBOT_IP:-192.168.123.164}"
ROBOT_USER="${ROBOT_USER:-unitree}"
ROBOT_ENV="${ROBOT_ENV:-marcus}"
REMOTE_DIR="/home/${ROBOT_USER}/Saqr"
SSH_OPTS="-o StrictHostKeyChecking=no -o ConnectTimeout=10"

RUN_AFTER=false

while [[ $# -gt 0 ]]; do
    case $1 in
        --run)    RUN_AFTER=true; shift ;;
        --ip)     ROBOT_IP="$2"; shift 2 ;;
        --env)    ROBOT_ENV="$2"; shift 2 ;;
        *)        echo "Unknown arg: $1"; exit 1 ;;
    esac
done

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

echo ""
echo "[1/4] 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
}

echo "[2/4] Creating remote directory..."
ssh ${SSH_OPTS} ${ROBOT_USER}@${ROBOT_IP} \
    "mkdir -p ${REMOTE_DIR}/runtime/{captures/{SAFE,PARTIAL,UNSAFE},logs,runs} ${REMOTE_DIR}/data/models"

echo "[3/4] Rsyncing package..."
RSYNC_OPTS=(
    -avz --delete
    --exclude="__pycache__" --exclude="*.pyc"
    --exclude=".git" --exclude="runtime" --exclude="data/dataset"
    --exclude="*.egg-info" --exclude="build/" --exclude="dist/"
)
rsync "${RSYNC_OPTS[@]}" \
    "${PROJECT_ROOT}/saqr" \
    "${PROJECT_ROOT}/scripts" \
    "${PROJECT_ROOT}/config" \
    "${PROJECT_ROOT}/docs" \
    "${PROJECT_ROOT}/pyproject.toml" \
    "${PROJECT_ROOT}/requirements.txt" \
    "${PROJECT_ROOT}/README.md" \
    ${ROBOT_USER}@${ROBOT_IP}:${REMOTE_DIR}/

# Weights (these are outside the package so rsync them separately).
if [ -f "${PROJECT_ROOT}/data/models/saqr_best.pt" ]; then
    echo "  Uploading saqr_best.pt..."
    scp ${SSH_OPTS} "${PROJECT_ROOT}/data/models/saqr_best.pt" \
        ${ROBOT_USER}@${ROBOT_IP}:${REMOTE_DIR}/data/models/
else
    echo "  [WARN] data/models/saqr_best.pt not found - train first!"
fi

echo "[4/4] Installing package on robot..."
ssh ${SSH_OPTS} ${ROBOT_USER}@${ROBOT_IP} bash -s <<INSTALL_EOF
    set -e
    source ~/miniconda3/etc/profile.d/conda.sh
    conda activate ${ROBOT_ENV}
    cd ${REMOTE_DIR}
    pip install -q -e . 2>/dev/null || pip install -e .
    chmod +x scripts/start_saqr.sh scripts/run_robot.sh scripts/run_local.sh 2>/dev/null || true
    echo "  Install OK"
INSTALL_EOF

if [ "$RUN_AFTER" = true ]; then
    echo ""
    echo "Starting Saqr bridge on robot..."
    ssh ${SSH_OPTS} -t ${ROBOT_USER}@${ROBOT_IP} bash -c "'${REMOTE_DIR}/scripts/start_saqr.sh'"
else
    echo ""
    echo "============================================"
    echo "  Deployed! Start the bridge via:"
    echo "============================================"
    echo ""
    echo "  ssh ${ROBOT_USER}@${ROBOT_IP}"
    echo "  ${REMOTE_DIR}/scripts/start_saqr.sh       # manual"
    echo "  sudo systemctl restart saqr-bridge        # systemd (see docs/start.md)"
    echo ""
fi