#!/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)"

config_get() {
    python3 -c "
import json, os, sys
with open('$PROJECT_ROOT/config/robot_config.json') as f:
    c = json.load(f)
for k in sys.argv[1].split('.'):
    c = c[k]
print(os.path.expandvars(str(c)))
" "$1"
}

ROBOT_IP="${ROBOT_IP:-$(config_get deploy.robot_ip)}"
ROBOT_USER="${ROBOT_USER:-$(config_get deploy.robot_user)}"
ROBOT_ENV="${ROBOT_ENV:-$(config_get deploy.robot_env)}"
REMOTE_DIR="${REMOTE_DIR:-$(config_get deploy.remote_dir)}"
SSH_OPTS="${SSH_OPTS:-$(config_get deploy.ssh_opts)}"

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} \
             ${REMOTE_DIR}/runtime/runs \
             ${REMOTE_DIR}/logs \
             ${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}/core" \
    "${PROJECT_ROOT}/apps" \
    "${PROJECT_ROOT}/gui" \
    "${PROJECT_ROOT}/robot" \
    "${PROJECT_ROOT}/utils" \
    "${PROJECT_ROOT}/scripts" \
    "${PROJECT_ROOT}/config" \
    "${PROJECT_ROOT}/docs" \
    "${PROJECT_ROOT}/assets" \
    "${PROJECT_ROOT}/pyproject.toml" \
    "${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 -e .
    chmod +x scripts/start_saqr.sh
    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