#!/usr/bin/env bash
# ──────────────────────────────────────────────────────────────────────────────
# Pi Multi-FX Pedal — DietPi first-boot post-install script
#
# Run AFTER a minimal DietPi install (no desktop, no DietPi-Software options).
# Installs only the packages needed for real-time guitar DSP with NAM.
#
# Usage:
#   sudo ./dietpi-postinstall.sh [--unattended]
#
# Environment:
#   PEDAL_USER:  system user to create (default: pi)
#   PEDAL_DIR:   installation directory (default: /opt/pi-multifx-pedal)
#
# Target footprint: < 60 MB idle RAM, < 500 MB disk
# ──────────────────────────────────────────────────────────────────────────────

set -euo pipefail

PEDAL_USER="${PEDAL_USER:-pi}"
PEDAL_DIR="${PEDAL_DIR:-/opt/pi-multifx-pedal}"
UNATTENDED=false
[[ "${1:-}" == "--unattended" ]] && UNATTENDED=true

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(dirname "$SCRIPT_DIR")"

# ── Colour helpers ───────────────────────────────────────────────────────────
info()  { printf "\e[34m[INFO] %s\e[0m\n" "$*"; }
ok()    { printf "\e[32m[ OK ] %s\e[0m\n" "$*"; }
warn()  { printf "\e[33m[WARN] %s\e[0m\n" "$*"; }
err()   { printf "\e[31m[FAIL] %s\e[0m\n" "$*"; }

# ── Root check ───────────────────────────────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
    err "This script must be run as root (sudo ./dietpi-postinstall.sh)"
    exit 1
fi

# ── Architecture check (must be RPi aarch64/armv7l) ──────────────────────────
ARCH="$(uname -m)"
if [[ "$ARCH" != "aarch64" && "$ARCH" != "armv7l" ]]; then
    err "This script is for Raspberry Pi (aarch64/armv7l), detected: $ARCH"
    exit 1
fi

info "========== Pi Multi-FX Pedal — DietPi Post-Install =========="
info "Architecture:   $ARCH"
info "Pedal user:     $PEDAL_USER"
info "Install dir:    $PEDAL_DIR"
echo ""

# ══════════════════════════════════════════════════════════════════════════════
# Step 1: System update & kernel
# ══════════════════════════════════════════════════════════════════════════════
info "Step 1: System update..."

apt-get update -qq
apt-get dist-upgrade -y -qq
ok "System updated"

# ──────────────────────────────────────────────────────────────────────────────
# Step 2: Install only the packages we actually need
# ──────────────────────────────────────────────────────────────────────────────
info "Step 2: Installing essential packages..."

# Core
apt-get install -y -qq \
    python3 \
    python3-pip \
    python3-venv \
    python3-dev \
    python3-numpy \
    git \
    ca-certificates \
    curl \
    rsync

# Audio
apt-get install -y -qq \
    jackd2 \
    alsa-utils \
    jack-tools \
    libjack-jackd2-dev

# Build tools for NAM dependencies (PyTorch, cmake, etc.)
apt-get install -y -qq \
    cmake \
    build-essential \
    pkg-config \
    libatlas-base-dev \
    libopenblas-dev

# Hardware interfaces
apt-get install -y -qq \
    device-tree-compiler \
    i2c-tools \
    raspi-gpio

ok "Essential packages installed"

# ──────────────────────────────────────────────────────────────────────────────
# Step 3: Install Python dependencies
# ──────────────────────────────────────────────────────────────────────────────
info "Step 3: Setting up Python virtual environment..."

mkdir -p "$PEDAL_DIR"
python3 -m venv "$PEDAL_DIR/.venv"
source "$PEDAL_DIR/.venv/bin/activate"

# Upgrade pip first
pip install --quiet --upgrade pip

# Install project dependencies from requirements.txt (if present locally)
if [[ -f "$PROJECT_DIR/requirements.txt" ]]; then
    pip install --quiet -r "$PROJECT_DIR/requirements.txt"
elif [[ -f "$PEDAL_DIR/requirements.txt" ]]; then
    pip install --quiet -r "$PEDAL_DIR/requirements.txt"
else
    # Manual install of known deps when requirements.txt not staged
    pip install --quiet \
        jackclient-python>=0.5.5 \
        numpy>=1.24 \
        neural-amp-modeler>=0.8.0 \
        scipy>=1.10 \
        RPi.GPIO>=0.7.1 \
        adafruit-circuitpython-ssd1306>=2.12 \
        adafruit-circuitpython-neopixel>=6.3 \
        pillow>=10.0 \
        python-rtmidi>=1.5 \
        pyserial>=3.5 \
        pyyaml>=6.0 \
        orjson>=3.9
fi

ok "Python virtual environment ready ($PEDAL_DIR/.venv)"

# ──────────────────────────────────────────────────────────────────────────────
# Step 4: Real-time audio limits
# ──────────────────────────────────────────────────────────────────────────────
info "Step 4: Setting real-time priority limits..."

LIMITS_FILE="/etc/security/limits.d/99-audio.conf"
cat > "$LIMITS_FILE" <<'LIMITS'
# Pi Multi-FX Pedal — real-time audio limits
@audio   -  rtprio     95
@audio   -  memlock    unlimited
@audio   -  nice       -20
LIMITS
ok "Created $LIMITS_FILE"

# Add pedal user to audio group
if ! groups "$PEDAL_USER" 2>/dev/null | grep -q '\baudio\b'; then
    usermod -a -G audio "$PEDAL_USER"
    info "Added '$PEDAL_USER' to audio group"
fi
# Also create 'pi' user if it exists on DietPi
if id pi &>/dev/null && ! groups pi | grep -q '\baudio\b'; then
    usermod -a -G audio pi
    info "Added 'pi' to audio group"
fi

ok "RT limits configured"

# ──────────────────────────────────────────────────────────────────────────────
# Step 5: Enable I2C (for OLED display)
# ──────────────────────────────────────────────────────────────────────────────
info "Step 5: Enabling I2C interface..."

CONFIG_TXT=""
if [[ -f /boot/firmware/config.txt ]]; then
    CONFIG_TXT="/boot/firmware/config.txt"
elif [[ -f /boot/config.txt ]]; then
    CONFIG_TXT="/boot/config.txt"
fi

if [[ -n "$CONFIG_TXT" ]]; then
    if ! grep -q "^dtparam=i2c_arm=on" "$CONFIG_TXT" 2>/dev/null; then
        echo "dtparam=i2c_arm=on" >> "$CONFIG_TXT"
        info "Enabled I2C in $CONFIG_TXT"
    else
        ok "I2C already enabled"
    fi
fi

# Load I2C kernel module now (so it works before reboot)
modprobe i2c_dev 2>/dev/null || true

ok "I2C configured"

# ──────────────────────────────────────────────────────────────────────────────
# Step 6: Disable unnecessary services (DietPi minimal baseline)
# ──────────────────────────────────────────────────────────────────────────────
info "Step 6: Disabling unnecessary services..."

SERVICES_TO_DISABLE=(
    bluetooth
    avahi-daemon
    triggerhappy
    ModemManager
    systemd-resolved
)

for svc in "${SERVICES_TO_DISABLE[@]}"; do
    if systemctl is-enabled "$svc" &>/dev/null 2>&1; then
        systemctl disable "$svc"
        systemctl stop "$svc" 2>/dev/null || true
        info "Disabled: $svc"
    fi
done

# PulseAudio conflicts with JACK
systemctl mask pulseaudio.service 2>/dev/null || true
systemctl mask pulseaudio.socket 2>/dev/null || true
info "Masked pulseaudio.service/socket"

# Optional: disable wpa_supplicant if wired-only (prompt on non-unattended)
if ! $UNATTENDED; then
    echo ""
    read -r -p "Is this pedal wired-only (disable WiFi)? [y/N] " WIRED_ONLY
    if [[ "$WIRED_ONLY" == "y" || "$WIRED_ONLY" == "Y" ]]; then
        systemctl disable wpa_supplicant 2>/dev/null || true
        systemctl mask wpa_supplicant 2>/dev/null || true
        info "Disabled wpa_supplicant (wired-only mode)"
    fi
fi

ok "Service minimization complete"

# ──────────────────────────────────────────────────────────────────────────────
# Step 7: Install pedal project files
# ──────────────────────────────────────────────────────────────────────────────
info "Step 7: Installing pedal project files..."

if [[ -d "$PROJECT_DIR/src" ]]; then
    rsync -a --delete \
        --exclude='.git' \
        --exclude='__pycache__' \
        --exclude='*.pyc' \
        --exclude='.venv' \
        --exclude='venv' \
        --exclude='hardware/gerber' \
        "$PROJECT_DIR/" "$PEDAL_DIR/"
    ok "Project files staged at $PEDAL_DIR"
else
    warn "No project source found at $PROJECT_DIR — skipping file copy"
    warn "Clone the repository first:"
    warn "  git clone https://gitea.ourpad.casa/shawn/pi-multifx-pedal.git $PEDAL_DIR"
fi

# ──────────────────────────────────────────────────────────────────────────────
# Step 8: Install systemd services
# ──────────────────────────────────────────────────────────────────────────────
info "Step 8: Installing systemd services..."

if [[ -f "$PEDAL_DIR/scripts/install_service.sh" ]]; then
    bash "$PEDAL_DIR/scripts/install_service.sh" "$PEDAL_DIR"
else
    warn "install_service.sh not found — skipping service installation"
fi

# ──────────────────────────────────────────────────────────────────────────────
# Step 9: Setup log rotation
# ──────────────────────────────────────────────────────────────────────────────
info "Step 9: Configuring log rotation..."

cat > /etc/logrotate.d/pi-multifx-pedal <<'LOGROTATE'
/var/log/pi-multifx-pedal/*.log {
    daily
    rotate 7
    compress
    delaycompress
    missingok
    notifempty
    copytruncate
}
LOGROTATE
ok "Log rotation configured"

# ──────────────────────────────────────────────────────────────────────────────
# Summary
# ──────────────────────────────────────────────────────────────────────────────
echo ""
info "========== Post-Install Complete =========="
echo ""
echo "  System user:    $PEDAL_USER"
echo "  Install dir:    $PEDAL_DIR"
echo "  Audio:          JACK + ALSA + RT limits"
echo "  I2C:            Enabled (for OLED)"
echo "  Unused svcs:    bluetooth, avahi, triggerhappy, ModemManager, resolved"
echo ""
echo "  Next steps:"
echo "    1. Run scripts/minimize.sh to strip further unused packages & kernel modules"
echo "    2. Run scripts/build-minimal-image.sh to produce a flashable image"
echo "    3. Reboot to verify everything comes up cleanly"
echo "       sudo reboot"
echo ""
echo "  Post-reboot verification:"
echo "    systemctl status pi-multifx-pedal.service"
echo "    systemctl status jackd.service"
echo "    free -m | grep Mem"
echo "    df -h /"
echo ""