#!/usr/bin/env bash
# ────────────────────────────────────────────────────────────────────
# Pi Multi-FX Pedal — First-boot audio setup
#
# Installs JACK + ALSA + I2S audio configuration for the RPi 4B.
# Must be run as root (sudo).
#
# Usage:
#   sudo ./setup_audio.sh [hat_type]
#
#   hat_type   Audio HAT key (default: audioinjector)
#              Options: audioinjector, pcm1808_pcm5102, iqaudio_codec,
#                       justboom, wm8731
# ────────────────────────────────────────────────────────────────────

set -euo pipefail

HAT_TYPE="${1:-audioinjector}"
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" "$*"; }

NEED_REBOOT=false

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

# ── Architecture check ─────────────────────────────────────────────
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 — Audio Setup =========="
info "HAT type:       $HAT_TYPE"
info "Architecture:   $ARCH"
echo ""

# ────────────────────────────────────────────────────────────────────
# Step 1: Install packages
# ────────────────────────────────────────────────────────────────────
info "Step 1: Installing audio packages..."

apt-get update -qq
apt-get install -y -qq \
    jackd2 \
    alsa-utils \
    jack-tools \
    jackd \
    libjack-jackd2-dev \
    python3-pip \
    2>&1 | tail -1

ok "Audio packages installed"

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

LIMITS_FILE="/etc/security/limits.d/99-audio.conf"
if [[ ! -f "$LIMITS_FILE" ]]; then
    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"
else
    ok "$LIMITS_FILE already exists"
fi

# Add 'pi' user to audio group if not already
if groups pi | grep -q '\baudio\b'; then
    ok "User 'pi' already in audio group"
else
    usermod -a -G audio pi
    info "Added 'pi' to audio group (will take effect on next login)"
fi

# ────────────────────────────────────────────────────────────────────
# Step 3: I2S overlay in config.txt
# ────────────────────────────────────────────────────────────────────
info "Step 3: Configuring I2S overlay in /boot/config.txt..."

declare -A OVERLAYS
OVERLAYS[audioinjector]="dtoverlay=audioinjector-wm8731"
OVERLAYS[pcm1808_pcm5102]="dtoverlay=audiosense-pi"
OVERLAYS[iqaudio_codec]="dtoverlay=iqaudio-codec"
OVERLAYS[justboom]="dtoverlay=justboom-dac"
OVERLAYS[wm8731]="dtoverlay=wm8731"

OVERLAY_LINE="${OVERLAYS[$HAT_TYPE]:-}"
if [[ -z "$OVERLAY_LINE" ]]; then
    err "Unknown HAT type '$HAT_TYPE'"
    echo "  Valid options: ${!OVERLAYS[*]}"
    exit 1
fi

CONFIG_TXT="/boot/config.txt"

if grep -qF "$OVERLAY_LINE" "$CONFIG_TXT"; then
    ok "I2S overlay already present: $OVERLAY_LINE"
else
    {
        echo ""
        echo "# Pi Multi-FX Pedal — $HAT_TYPE"
        echo "$OVERLAY_LINE"
    } >> "$CONFIG_TXT"
    ok "Appended to $CONFIG_TXT: $OVERLAY_LINE"
    NEED_REBOOT=true
fi

# Enable I2C (for OLED / sensor I2C) if not already
if ! grep -q "^dtparam=i2c_arm=on" "$CONFIG_TXT"; then
    echo "dtparam=i2c_arm=on" >> "$CONFIG_TXT"
    info "Enabled I2C interface"
fi

# ────────────────────────────────────────────────────────────────────
# Step 4: Disable onboard audio (headphone jack) to avoid conflicts
# ────────────────────────────────────────────────────────────────────
info "Step 4: Disabling onboard analog audio (headphone jack)..."

if grep -q "^dtparam=audio=on" "$CONFIG_TXT"; then
    sed -i 's/^dtparam=audio=on/dtparam=audio=off/' "$CONFIG_TXT"
    info "Disabled onboard analog audio (dtparam=audio=off)"
    NEED_REBOOT=true
elif grep -q "^dtparam=audio=off" "$CONFIG_TXT"; then
    ok "Onboard audio already disabled"
else
    echo "dtparam=audio=off" >> "$CONFIG_TXT"
    info "Disabled onboard analog audio"
fi

# ────────────────────────────────────────────────────────────────────
# Step 5: Systemd service for JACK
# ────────────────────────────────────────────────────────────────────
info "Step 5: Installing JACK systemd service..."

SERVICE_FILE="/etc/systemd/system/jackd.service"
PYTHON_BIN="/usr/bin/python3"
GENERATE_SCRIPT=$(cat <<'PYEOF'
import sys
sys.path.insert(0, '/home/pi/pi-multifx-pedal/src')
from system.audio import AudioConfig, AudioSystem
hat = sys.argv[1] if len(sys.argv) > 1 else "audioinjector"
cfg = AudioConfig(hat_type=hat)
print(AudioSystem.systemd_service_content(cfg))
PYEOF
)

# Dynamic generation via Python so service stays in sync with code
if [[ -d "$PROJECT_DIR/src" ]]; then
    PYTHONPATH="$PROJECT_DIR/src" python3 -c "
import sys
sys.path.insert(0, '$PROJECT_DIR/src')
from system.audio import AudioConfig, AudioSystem
cfg = AudioConfig(hat_type='$HAT_TYPE')
print(AudioSystem.systemd_service_content(cfg))
" > "$SERVICE_FILE"
    ok "Generated $SERVICE_FILE from AudioSystem.systemd_service_content()"
else
    # Fallback: hardcoded unit
    cat > "$SERVICE_FILE" <<UNIT
[Unit]
Description=JACK Audio Server — Pi Multi-FX Pedal
After=sound.target network.target
Wants=multi-fx-pedal.target

[Service]
Type=simple
User=pi
ExecStart=/usr/bin/jackd -P70 -p128 -n2 -r48000 -dalsa -dhw:0 -i2 -o2
Restart=on-failure
RestartSec=5
LimitRTPRIO=95
LimitMEMLOCK=infinity

[Install]
WantedBy=multi-user.target
UNIT
    ok "Created fallback $SERVICE_FILE (src directory not found)"
fi

# ────────────────────────────────────────────────────────────────────
# Step 6: Silence ALSA mixer defaults
# ────────────────────────────────────────────────────────────────────
info "Step 6: Setting default ALSA volumes..."

# Set a reasonable default capture gain for guitar input
if command -v amixer &>/dev/null; then
    # Try to unmute and set capture volume on card 0
    amixer -c 0 set 'Mic' 80% unmute 2>/dev/null || true
    amixer -c 0 set 'Capture' 80% unmute 2>/dev/null || true
    ok "ALSA mixer defaults set"
else
    info "amixer not available — skipping volume config"
fi

# ────────────────────────────────────────────────────────────────────
# Step 7: Systemd reload & enable
# ────────────────────────────────────────────────────────────────────
info "Step 7: Reloading systemd and enabling JACK service..."

systemctl daemon-reload
systemctl enable jackd.service
ok "JACK service enabled (start on next boot)"

# ────────────────────────────────────────────────────────────────────
# Summary
# ────────────────────────────────────────────────────────────────────
echo ""
info "========== Setup Complete =========="

if [[ "${NEED_REBOOT:-false}" == true ]]; then
    warn "REBOOT REQUIRED — I2S overlay changes need a reboot to take effect"
    warn "Run: sudo reboot"
else
    ok "No reboot required — you can start JACK now:"
    echo "    sudo systemctl start jackd"
fi

echo ""
info "Post-setup verification:"
echo "    1. Check ALSA devices:        aplay -l && arecord -l"
echo "    2. Verify JACK status:        jack_wait -c"
echo "    3. Measure latency:           jack_iodelay"
echo "    4. Listen for xruns:          jack_showtime -c"
echo "    5. Run Python test:           python3 -c "
echo "                                   \"from system.audio import AudioSystem; "
echo "                                    sys = AudioSystem(); "
echo "                                    print(sys.list_devices())\""
echo ""
info "Documented ALSA device names:"
echo "    Card 0 — I2S HAT ($HAT_TYPE)"
echo "      Playback:  hw:0,0  (DAC output)"
echo "      Capture:   hw:0,0  (ADC input)"
echo "    Card 1 — (if present) USB audio / HDMI"