#!/usr/bin/env bash
# Robot setup for Raspberry Pi OS Lite (64-bit) on a Pi Zero 2 W.
# Run as your normal user (not root) from inside this "pi" folder:
#
#     ./install.sh <omarchy-machine-ip> [--servos]
#
# Safe to re-run. Reboot afterwards so the I2S sound card overlay takes effect.
set -euo pipefail

HOST_IP="${1:-}"
SERVOS=0
[[ "${2:-}" == "--servos" ]] && SERVOS=1

if [[ -z "$HOST_IP" || ! "$HOST_IP" =~ ^[0-9]{1,3}(\.[0-9]{1,3}){3}$ ]]; then
    echo "usage: $0 <omarchy-machine-ipv4> [--servos]" >&2
    exit 1
fi
if [[ $EUID -eq 0 ]]; then
    echo "Run this as your normal user; it calls sudo where needed." >&2
    exit 1
fi

HERE="$(cd "$(dirname "$0")" && pwd)"
ME="$(id -un)"
CONFIG_TXT=/boot/firmware/config.txt
[[ -f $CONFIG_TXT ]] || CONFIG_TXT=/boot/config.txt

echo "==> Installing packages"
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
    python3-picamera2 python3-venv \
    pipewire pipewire-pulse pipewire-alsa wireplumber libspa-0.2-modules \
    pulseaudio-utils alsa-utils nftables openssl

echo "==> Enabling the I2S microphone/amplifier overlay in $CONFIG_TXT"
add_line() { grep -qxF "$1" "$CONFIG_TXT" || echo "$1" | sudo tee -a "$CONFIG_TXT" >/dev/null; }
add_line "dtparam=i2s=on"
add_line "dtoverlay=googlevoicehat-soundcard"
if [[ $SERVOS -eq 1 ]]; then
    add_line "dtparam=i2c_arm=on"
fi

echo "==> Installing the MCP server to /opt/robot-mcp"
sudo mkdir -p /opt/robot-mcp
sudo chown "$ME:$ME" /opt/robot-mcp
cp "$HERE/robot_mcp_server.py" /opt/robot-mcp/
# --system-site-packages so the venv can see the apt-installed picamera2/libcamera
[[ -d /opt/robot-mcp/venv ]] || python3 -m venv --system-site-packages /opt/robot-mcp/venv
/opt/robot-mcp/venv/bin/pip install --upgrade pip >/dev/null
/opt/robot-mcp/venv/bin/pip install -r "$HERE/requirements.txt"
if [[ $SERVOS -eq 1 ]]; then
    /opt/robot-mcp/venv/bin/pip install adafruit-circuitpython-servokit
fi

echo "==> Writing /etc/robot-mcp.env"
if ! sudo test -f /etc/robot-mcp.env; then
    TOKEN="$(openssl rand -hex 32)"
    sudo tee /etc/robot-mcp.env >/dev/null <<EOF
ROBOT_MCP_TOKEN=$TOKEN
ROBOT_MCP_PORT=8765
ROBOT_CAM_SIZE=1280x720
ROBOT_CAM_HFLIP=0
ROBOT_CAM_VFLIP=0
ROBOT_SERVOS=$SERVOS
EOF
    sudo chown root:"$ME" /etc/robot-mcp.env
    sudo chmod 640 /etc/robot-mcp.env
else
    sudo sed -i "s/^ROBOT_SERVOS=.*/ROBOT_SERVOS=$SERVOS/" /etc/robot-mcp.env
fi

echo "==> Installing the systemd service"
sed "s/__USER__/$ME/g" "$HERE/robot-mcp.service" | sudo tee /etc/systemd/system/robot-mcp.service >/dev/null
getent group i2c >/dev/null || sudo groupadd --system i2c
sudo systemctl daemon-reload
sudo systemctl enable robot-mcp.service

echo "==> Configuring PipeWire to serve audio to $HOST_IP"
mkdir -p ~/.config/pipewire/pipewire-pulse.conf.d
cp "$HERE/pipewire/50-network-server.conf" ~/.config/pipewire/pipewire-pulse.conf.d/
# Run the audio stack at boot without anyone logging in.
sudo loginctl enable-linger "$ME"
systemctl --user enable pipewire.socket pipewire-pulse.socket wireplumber.service 2>/dev/null || true

echo "==> Restricting ports 4713 and 8765 to $HOST_IP"
sed "s/__HOST_IP__/$HOST_IP/g" "$HERE/nftables-robot.nft" | sudo tee /etc/nftables-robot.nft >/dev/null
grep -qxF 'include "/etc/nftables-robot.nft"' /etc/nftables.conf \
    || echo 'include "/etc/nftables-robot.nft"' | sudo tee -a /etc/nftables.conf >/dev/null
sudo systemctl enable nftables.service
sudo systemctl restart nftables.service

echo
echo "Done. Reboot now:  sudo reboot"
echo
echo "Your MCP token (put this in ~/.hermes/.env on the Omarchy machine as ROBOT_MCP_TOKEN):"
sudo grep '^ROBOT_MCP_TOKEN=' /etc/robot-mcp.env | cut -d= -f2