#!/usr/bin/env bash
# Wii Guitar -> USB MIDI bridge : installer for a Raspberry Pi Zero 2 W.
#
#   sudo ./install.sh              install
#   sudo ./install.sh --dry-run    show what it would do, change nothing
#   sudo ./install.sh --user pi    install for a specific user
#
# Safe to re-run: every step checks before it changes anything.

set -uo pipefail

VERSION="1.0"
DRY=0
TARGET_USER=""
SRC_URL="https://jacksonmicro.com/wiiguitar/wiiguitar.tar.gz"

while [ $# -gt 0 ]; do
    case "$1" in
        --dry-run) DRY=1 ;;
        --user) TARGET_USER="${2:-}"; shift ;;
        -h|--help) sed -n '2,10p' "$0" | sed 's/^# \?//'; exit 0 ;;
        *) echo "unknown option: $1" >&2; exit 2 ;;
    esac
    shift
done

bold()  { printf '\033[1m%s\033[0m\n' "$*"; }
info()  { printf '  %s\n' "$*"; }
warn()  { printf '\033[33m  %s\033[0m\n' "$*"; }
die()   { printf '\033[31m%s\033[0m\n' "$*" >&2; exit 1; }
run()   { if [ "$DRY" = 1 ]; then printf '  [dry-run] %s\n' "$*"; else eval "$@"; fi; }

[ "$DRY" = 1 ] || [ "$(id -u)" = 0 ] || die "run with sudo:  sudo $0 $*"

# ---------------------------------------------------------------- user ----
if [ -z "$TARGET_USER" ]; then
    TARGET_USER="${SUDO_USER:-}"
    if [ -z "$TARGET_USER" ] || [ "$TARGET_USER" = root ]; then
        for u in dietpi pi; do
            if id "$u" >/dev/null 2>&1; then TARGET_USER="$u"; break; fi
        done
    fi
fi
id "$TARGET_USER" >/dev/null 2>&1 || die "no such user: '$TARGET_USER' (use --user)"
HOME_DIR="$(getent passwd "$TARGET_USER" | cut -d: -f6)"
APP_DIR="$HOME_DIR/wiiguitar"

bold "Wii Guitar bridge installer v$VERSION"
info "user      : $TARGET_USER"
info "install to: $APP_DIR"
[ "$DRY" = 1 ] && warn "DRY RUN - nothing will be changed"

# -------------------------------------------------------------- checks ----
MODEL="$(tr -d '\0' < /proc/device-tree/model 2>/dev/null)"; MODEL="${MODEL:-unknown}"
info "board     : $MODEL"
case "$MODEL" in
    *"Zero 2"*) ;;
    *Raspberry*) warn "Not a Zero 2 W. USB gadget mode needs a Pi with USB OTG"
                 warn "(Zero / Zero 2 / A-series / CM). A 3B or 4B model B will not work." ;;
    *) warn "Not a Raspberry Pi? The USB MIDI gadget needs a USB device controller." ;;
esac

# ------------------------------------------------------------ packages ----
bold "1/7  packages"
PKGS="python3 python3-numpy python3-pil python3-spidev python3-evdev
      python3-libgpiod python3-lgpio bluez bluez-firmware alsa-utils
      fonts-dejavu-core dbus"
# pi-bluetooth ships the Pi's BT helpers; absent on non-Pi Debian.
apt-cache show pi-bluetooth >/dev/null 2>&1 && PKGS="$PKGS pi-bluetooth"
MISSING=""
for p in $PKGS; do
    dpkg -s "$p" >/dev/null 2>&1 || MISSING="$MISSING $p"
done
if [ -n "$MISSING" ]; then
    info "installing:$MISSING"
    run "apt-get update -qq"
    run "DEBIAN_FRONTEND=noninteractive apt-get install -y -qq$MISSING"
else
    info "all present"
fi

# ---------------------------------------------------------------- boot ----
bold "2/7  boot config (SPI + USB gadget)"
# WIIGUITAR_BOOT lets the test suite point this at a sandbox.
BOOT="${WIIGUITAR_BOOT:-}"
if [ -z "$BOOT" ]; then
    BOOT=/boot/firmware
    [ -d "$BOOT" ] || BOOT=/boot
fi
CFG="$BOOT/config.txt"; CMD="$BOOT/cmdline.txt"
[ -f "$CFG" ] || die "no $CFG - is this a Raspberry Pi?"

REBOOT_NEEDED=0
if ! grep -q '^dtparam=spi=on' "$CFG"; then
    info "enabling SPI (the LCD HAT)"
    run "cp -n '$CFG' '$CFG.bak-wiiguitar'"
    run "sed -i 's/^#\\?dtparam=spi=.*/dtparam=spi=on/' '$CFG'"
    grep -q '^dtparam=spi=on' "$CFG" || run "printf '\\ndtparam=spi=on\\n' >> '$CFG'"
    REBOOT_NEEDED=1
else
    info "SPI already on"
fi

# dr_mode=peripheral: the Pi must be the USB *device*; the phone/PC is host.
if ! grep -q 'dtoverlay=dwc2' "$CFG"; then
    info "enabling USB gadget mode (dwc2, peripheral)"
    run "printf '\\n# Wii Guitar USB MIDI gadget\\ndtoverlay=dwc2,dr_mode=peripheral\\n' >> '$CFG'"
    REBOOT_NEEDED=1
else
    info "dwc2 overlay already set"
fi

if [ -f "$CMD" ] && ! grep -q 'modules-load=dwc2' "$CMD"; then
    info "loading dwc2 at boot"
    run "cp -n '$CMD' '$CMD.bak-wiiguitar'"
    run "sed -i 's/rootwait/rootwait modules-load=dwc2/' '$CMD'"
    REBOOT_NEEDED=1
fi

# ----------------------------------------------------------- bluetooth ----
bold "3/7  Bluetooth (the parts that make a Wii Remote work)"
# A Wii Remote paired with 1+2 never bonds: its PIN is a raw binary value the
# standard BlueZ agent cannot supply. BlueZ >=5.65 then refuses the HID
# connection ("Rejected connection from !bonded device") and no input device
# appears even though the controller reads as connected.
if ! grep -qs 'ClassicBondedOnly=false' /etc/bluetooth/input.conf 2>/dev/null; then
    info "allowing non-bonded HID connections"
    run "mkdir -p /etc/bluetooth"
    run "cat > /etc/bluetooth/input.conf <<'EOF'
[General]
ClassicBondedOnly=false
UserspaceHID=true
IdleTimeout=0
EOF"
else
    info "input.conf already set"
fi

# Wiimotes need ERTM disabled or the L2CAP connection is refused.
if ! grep -qs 'disable_ertm=1' /etc/modprobe.d/bluetooth-wiimote.conf 2>/dev/null; then
    info "disabling Bluetooth ERTM"
    run "echo 'options bluetooth disable_ertm=1' > /etc/modprobe.d/bluetooth-wiimote.conf"
    REBOOT_NEEDED=1
else
    info "ERTM already disabled"
fi

run "printf 'hidp\\nuhid\\nhid-wiimote\\n' > /etc/modules-load.d/wiiguitar.conf"
for m in uhid hid-wiimote; do run "modprobe $m 2>/dev/null || true"; done
run "systemctl enable --now bluetooth >/dev/null 2>&1 || true"

# ----------------------------------------------------------- app files ----
bold "4/7  application"
HERE="$(cd "$(dirname "$0")" && pwd)"
FILES="chords.py modes.py songs.py backends.py hat.py ui.py bridge.py
       wiiguitar.py discover.py axistest.py pair.sh README.md"
if [ -f "$HERE/bridge.py" ]; then
    info "installing from $HERE"
    run "install -d -o '$TARGET_USER' -g '$TARGET_USER' '$APP_DIR'"
    for f in $FILES; do
        [ -f "$HERE/$f" ] && run "install -o '$TARGET_USER' -g '$TARGET_USER' -m 644 '$HERE/$f' '$APP_DIR/$f'"
    done
    run "chmod +x '$APP_DIR'/*.sh '$APP_DIR/bridge.py' '$APP_DIR/discover.py' 2>/dev/null || true"
else
    info "sources not alongside installer; fetching $SRC_URL"
    TMP="$(mktemp -d)"
    run "curl -fsSL '$SRC_URL' -o '$TMP/w.tar.gz'" || die "download failed"
    run "tar xzf '$TMP/w.tar.gz' -C '$TMP'"
    run "install -d -o '$TARGET_USER' -g '$TARGET_USER' '$APP_DIR'"
    run "cp -r '$TMP'/wiiguitar/* '$APP_DIR'/"
    run "chown -R '$TARGET_USER:$TARGET_USER' '$APP_DIR'"
    run "rm -rf '$TMP'"
fi

# --------------------------------------------------------- permissions ----
bold "5/7  permissions"
for g in gpio spi audio input; do
    getent group "$g" >/dev/null 2>&1 || run "groupadd -f '$g'"
done
if id -nG "$TARGET_USER" | tr ' ' '\n' | grep -qx gpio; then
    info "$TARGET_USER already in the hardware groups"
else
    info "adding $TARGET_USER to gpio, spi, audio, input"
    run "usermod -aG gpio,spi,audio,input '$TARGET_USER'"
fi

# ------------------------------------------------------------ services ----
bold "6/7  services"
run "cat > /usr/local/sbin/wiiguitar-gadget.sh <<'EOF'
#!/bin/bash
# Present this Pi to a USB host as a class-compliant MIDI device.
set -e
G=/sys/kernel/config/usb_gadget/wiiguitar
modprobe libcomposite
[ -d /sys/kernel/config/usb_gadget ] || mount -t configfs none /sys/kernel/config
if [ -d \"\$G\" ] && [ -s \"\$G/UDC\" ]; then exit 0; fi
mkdir -p \"\$G\"; cd \"\$G\"
echo 0x1d6b > idVendor
echo 0x0104 > idProduct
echo 0x0100 > bcdDevice
echo 0x0200 > bcdUSB
mkdir -p strings/0x409
echo \"\$(cat /sys/class/net/wlan0/address 2>/dev/null | tr -d ':' || echo 000000)\" > strings/0x409/serialnumber
echo \"WiiGuitar\"              > strings/0x409/manufacturer
echo \"Wii Guitar MIDI Bridge\" > strings/0x409/product
mkdir -p configs/c.1/strings/0x409
echo \"MIDI\" > configs/c.1/strings/0x409/configuration
# iOS refuses to enumerate accessories that ask for too much current.
echo 100 > configs/c.1/MaxPower
mkdir -p functions/midi.usb0
echo 1         > functions/midi.usb0/in_ports
echo 1         > functions/midi.usb0/out_ports
echo WiiGuitar > functions/midi.usb0/id
echo 128       > functions/midi.usb0/buflen
echo 32        > functions/midi.usb0/qlen
ln -sf functions/midi.usb0 configs/c.1/
ls /sys/class/udc > UDC
EOF"
run "chmod +x /usr/local/sbin/wiiguitar-gadget.sh"

run "cat > /usr/local/sbin/wiiguitar-connect.sh <<'EOF'
#!/bin/bash
# A Wii Remote paired with 1+2 never bonds, so it cannot re-initiate a
# connection itself: the host has to. Retry whenever the guitar is absent, so
# pressing 1+2 is all the user ever needs to do.
LAST_SCAN=0
SCAN_EVERY=30
have_guitar() { grep -qil guitar /sys/class/input/event*/device/name 2>/dev/null; }
nintendo_macs() {
    bluetoothctl devices 2>/dev/null | grep -iE 'Nintendo|RVL-CNT' | awk '{print \$2}'
}
while true; do
    if have_guitar; then sleep 5; continue; fi
    for mac in \$(nintendo_macs); do
        timeout 12 bluetoothctl connect \"\$mac\" >/dev/null 2>&1
        have_guitar && break
    done
    if ! have_guitar && [ \$(( \$(date +%s) - LAST_SCAN )) -ge \$SCAN_EVERY ]; then
        LAST_SCAN=\$(date +%s)
        timeout 12 bluetoothctl --timeout 8 scan on >/dev/null 2>&1
        for mac in \$(nintendo_macs); do
            bluetoothctl trust \"\$mac\" >/dev/null 2>&1
            timeout 12 bluetoothctl connect \"\$mac\" >/dev/null 2>&1
            have_guitar && break
        done
    fi
    sleep 5
done
EOF"
run "chmod +x /usr/local/sbin/wiiguitar-connect.sh"

run "cat > /etc/systemd/system/wiiguitar-gadget.service <<'EOF'
[Unit]
Description=Wii Guitar USB MIDI gadget
After=systemd-modules-load.service
[Service]
Type=oneshot
RemainAfterExit=yes
ExecStart=/usr/local/sbin/wiiguitar-gadget.sh
[Install]
WantedBy=multi-user.target
EOF"

run "cat > /etc/systemd/system/wiiguitar-connect.service <<'EOF'
[Unit]
Description=Keep the Wii Remote connected
After=bluetooth.service
Wants=bluetooth.service
[Service]
Type=simple
ExecStart=/usr/local/sbin/wiiguitar-connect.sh
Restart=always
RestartSec=5
Nice=10
[Install]
WantedBy=multi-user.target
EOF"

run "cat > /etc/systemd/system/wiiguitar.service <<EOF
[Unit]
Description=Wii Guitar -> USB MIDI bridge
After=bluetooth.service wiiguitar-gadget.service
Wants=bluetooth.service wiiguitar-gadget.service

[Service]
Type=simple
User=$TARGET_USER
WorkingDirectory=$APP_DIR
ExecStart=/usr/bin/python3 $APP_DIR/bridge.py --backend midi
Restart=always
RestartSec=3
Nice=-5
IOSchedulingClass=realtime

[Install]
WantedBy=multi-user.target
EOF"

run "systemctl daemon-reload"
run "systemctl enable wiiguitar-gadget wiiguitar-connect wiiguitar >/dev/null 2>&1"

# ---------------------------------------------------------------- done ----
bold "7/7  done"
if [ "$DRY" = 1 ]; then
    echo
    bold "Dry run complete - nothing was changed."
    exit 0
fi
echo
if [ "$REBOOT_NEEDED" = 1 ]; then
    bold "Reboot required (SPI / USB gadget / ERTM are set at boot):"
    echo "    sudo systemctl start reboot.target    # DietPi masks logind, so"
    echo "                                          # 'systemctl reboot' may fail"
else
    run "systemctl restart wiiguitar-gadget wiiguitar-connect wiiguitar"
fi
echo
bold "Then:"
echo "    cd $APP_DIR && sudo ./pair.sh     # press 1+2 on the Wii Remote"
echo "    plug the OUTER micro-USB port into your phone or computer"
echo
echo "  Check on it with:  journalctl -u wiiguitar -f"