#!/usr/bin/env bash
# Jetson Nano host setup script
# Run once on fresh JetPack 4.6 installation
# Usage: sudo bash setup-jetson.sh

set -euo pipefail

echo "=== Jetson Nano Host Setup — saltybot ==="
echo "JetPack 4.6 / L4T R32.6.1 expected"

# ── Verify we're on Jetson ────────────────────────────────────────────────────
if ! uname -m | grep -q aarch64; then
    echo "ERROR: Must run on Jetson (aarch64). Got: $(uname -m)"
    exit 1
fi

# ── System update ─────────────────────────────────────────────────────────────
apt-get update && apt-get upgrade -y

# ── Install Docker + NVIDIA runtime ──────────────────────────────────────────
if ! command -v docker &>/dev/null; then
    echo "[+] Installing Docker..."
    curl -fsSL https://get.docker.com | sh
    usermod -aG docker "$SUDO_USER"
fi

# NVIDIA container runtime
if ! dpkg -l | grep -q nvidia-container-runtime; then
    echo "[+] Installing NVIDIA Container Runtime..."
    distribution=$(. /etc/os-release; echo $ID$VERSION_ID)
    curl -s -L https://nvidia.github.io/libnvidia-container/gpgkey | apt-key add -
    curl -s -L "https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list" \
        > /etc/apt/sources.list.d/nvidia-container-runtime.list
    apt-get update
    apt-get install -y nvidia-container-runtime
fi

# Configure Docker daemon for NVIDIA runtime
cat > /etc/docker/daemon.json << 'EOF'
{
    "runtimes": {
        "nvidia": {
            "path": "nvidia-container-runtime",
            "runtimeArgs": []
        }
    },
    "default-runtime": "nvidia"
}
EOF
systemctl restart docker

# ── Set Jetson power mode to MAXN 10W ────────────────────────────────────────
echo "[+] Setting MAXN 10W power mode..."
nvpmodel -m 0
jetson_clocks

# ── Install udev rules ────────────────────────────────────────────────────────
echo "[+] Installing udev rules..."
cat > /etc/udev/rules.d/99-saltybot.rules << 'EOF'
# RPLIDAR A1M8 (SiliconLabs CP2102)
KERNEL=="ttyUSB*", ATTRS{idVendor}=="10c4", ATTRS{idProduct}=="ea60", \
  SYMLINK+="rplidar", MODE="0666"

# STM32 USB CDC (STMicroelectronics Virtual COM)
KERNEL=="ttyACM*", ATTRS{idVendor}=="0483", ATTRS{idProduct}=="5740", \
  SYMLINK+="stm32-bridge", MODE="0666"

# Intel RealSense D435i
SUBSYSTEM=="usb", ATTRS{idVendor}=="8086", ATTRS{idProduct}=="0b3a", \
  MODE="0666"
EOF

udevadm control --reload-rules
udevadm trigger

# ── Install RealSense udev rules ──────────────────────────────────────────────
echo "[+] Installing RealSense udev rules..."
if [ -f /etc/udev/rules.d/99-realsense-libusb.rules ]; then
    echo "  Already installed."
else
    # Download from librealsense repo
    wget -q -O /etc/udev/rules.d/99-realsense-libusb.rules \
        https://raw.githubusercontent.com/IntelRealSense/librealsense/master/config/99-realsense-libusb.rules
    udevadm control --reload-rules
    udevadm trigger
fi

# ── Enable I2C + UART ─────────────────────────────────────────────────────────
echo "[+] Enabling I2C and UART..."
modprobe i2c-dev
# Add user to i2c and dialout groups
usermod -aG i2c,dialout,gpio "$SUDO_USER"

# ── Configure UART (disable console on ttyTHS1) ───────────────────────────────
# ttyTHS1 is the 40-pin header UART — disable serial console to free it
if grep -q "console=ttyS0" /boot/extlinux/extlinux.conf; then
    echo "[+] UART ttyTHS1 already free for application use."
else
    echo "[!] Warning: Check /boot/extlinux/extlinux.conf if serial console"
    echo "    is using ttyTHS1. Disable it to use UART for STM32 bridge."
fi

# ── Docker Compose ────────────────────────────────────────────────────────────
if ! command -v docker-compose &>/dev/null && ! docker compose version &>/dev/null 2>&1; then
    echo "[+] Installing docker-compose..."
    pip3 install docker-compose
fi

# ── Swap (improve stability under memory pressure) ────────────────────────────
if [ "$(swapon --show | wc -l)" -le 1 ]; then
    echo "[+] Creating 4GB swap file..."
    fallocate -l 4G /swapfile
    chmod 600 /swapfile
    mkswap /swapfile
    swapon /swapfile
    echo '/swapfile none swap sw 0 0' >> /etc/fstab
fi

echo ""
echo "=== Setup complete ==="
echo "Please log out and back in for group membership to take effect."
echo ""
echo "Next steps:"
echo "  1. cd jetson/"
echo "  2. docker compose build"
echo "  3. docker compose up -d"
echo "  4. docker compose logs -f"