#!/bin/bash
# ──────────────────────────────────────────────────────────────────
# install_hat.sh — Enable PCM1808+PCM5102 I2S HAT overlay + test
# Pi Multi-FX Pedal project
# ──────────────────────────────────────────────────────────────────
set -euo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
PROJECT_ROOT="$(dirname "$HERE")"
OVERLAY_SRC="${PROJECT_ROOT}/hardware/pcm1808-pcm5102-overlay.dts"

# ── Config paths ──────────────────────────────────────────────
if [ -f /boot/firmware/config.txt ]; then
    CONFIG_FILE="/boot/firmware/config.txt"
    OVERLAY_DIR="/boot/firmware/overlays"
elif [ -f /boot/config.txt ]; then
    CONFIG_FILE="/boot/config.txt"
    OVERLAY_DIR="/boot/overlays"
else
    echo "ERROR: Cannot find config.txt"
    exit 1
fi

OVERLAY_DST="${OVERLAY_DIR}/pcm1808-pcm5102-overlay.dtbo"

# ── Color helpers ─────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'
info()  { echo -e "${GREEN}[✓]${NC} $1"; }
warn()  { echo -e "${YELLOW}[!]${NC} $1"; }
err()   { echo -e "${RED}[✗]${NC} $1"; }

# ── Step 1: Check prerequisites ──────────────────────────────
echo "━━━━━ PCM1808+PCM5102 I2S HAT Installer ━━━━━"
echo ""

if [ ! -f "$OVERLAY_SRC" ]; then
    err "Overlay source not found: $OVERLAY_SRC"
    exit 1
fi

if ! command -v dtc &>/dev/null; then
    echo "Installing device-tree-compiler..."
    sudo apt-get update -qq && sudo apt-get install -y -qq device-tree-compiler
fi

if ! command -v dtc &>/dev/null; then
    err "Failed to install device-tree-compiler"
    exit 1
fi
info "device-tree-compiler available"

# ── Step 2: Compile overlay ──────────────────────────────────
echo ""
echo "Compiling overlay..."
sudo dtc -@ -I dts -O dtb -o "$OVERLAY_DST" "$OVERLAY_SRC"
info "Overlay compiled → $OVERLAY_DST"

# ── Step 3: Enable overlay in config.txt ─────────────────────
echo ""
echo "Enabling overlay in $CONFIG_FILE..."

# Remove any existing pcm1808 overlay lines
sudo sed -i '/dtoverlay=pcm1808-pcm5102-overlay/d' "$CONFIG_FILE" 2>/dev/null || true
sudo sed -i '/dtparam=audio=off/d' "$CONFIG_FILE" 2>/dev/null || true
sudo sed -i '/dtparam=i2s=on/d' "$CONFIG_FILE" 2>/dev/null || true

# Add at the end
{
    echo ""
    echo "# ── Custom I2S HAT: PCM1808 ADC + PCM5102 DAC ──"
    echo "dtoverlay=pcm1808-pcm5102-overlay"
    echo "dtparam=audio=off"
} | sudo tee -a "$CONFIG_FILE" > /dev/null

info "Overlay added to $CONFIG_FILE"

# ── Step 4: Create ALSA config ───────────────────────────────
echo ""
echo "Setting up ALSA configuration..."

sudo bash -c "cat > /etc/asound.conf" << 'EOCONF'
# PCM1808+PCM5102 I2S HAT — default sound config
pcm.!default {
    type asym
    playback.pcm { type hw card 0 device 0 }
    capture.pcm  { type hw card 0 device 0 }
}
ctl.!default {
    type hw card 0
}
EOCONF

info "ALSA config written to /etc/asound.conf"

# ── Step 5: Disable onboard audio ────────────────────────────
echo ""
echo "Disabling onboard audio (PWM)..."
sudo sed -i 's/^dtparam=audio=on/# dtparam=audio=on/' "$CONFIG_FILE" 2>/dev/null || true
sudo sed -i 's/^#dtparam=audio=off/dtparam=audio=off/' "$CONFIG_FILE" 2>/dev/null || true

# Also mask the onboard audio kernel module
if lsmod | grep -q "snd_bcm2835"; then
    warn "Onboard audio module loaded — will be disabled after reboot"
fi

# ── Step 6: Disable PulseAudio (conflicts with JACK) ─────────
echo ""
echo "Masking PulseAudio for audio profile..."
sudo systemctl mask pulseaudio.service 2>/dev/null || true
sudo systemctl mask pulseaudio.socket 2>/dev/null || true
info "PulseAudio masked (JACK will take over)"

# ── Step 7: Reboot prompt ────────────────────────────────────
echo ""
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
info "Installation complete!"
echo ""
echo "Reboot required to load the overlay."
echo ""
echo "After reboot, run the test suite:"
echo "  $0 --test"
echo ""
echo "Or test manually:"
echo "  aplay -l        # should show pcm1808pcm5102hat"
echo "  arecord -l      # should show capture device"
echo "  speaker-test -t sine -f 440 -l 1"
echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

# ── Test mode (--test flag) ──────────────────────────────────
if [ "${1:-}" = "--test" ]; then
    echo ""
    echo "━━━━━ Running post-install tests ━━━━━"
    FAILED=0

    # Test 1: Overlay loaded
    echo -n "Test 1: Overlay loaded... "
    if dmesg | grep -qi "pcm1808-pcm5102"; then
        info "OK (found in dmesg)"
    else
        # Try via /proc
        if grep -q "pcm1808-pcm5102" /proc/device-tree/sound/name 2>/dev/null; then
            info "OK (found in device tree)"
        else
            err "FAILED — overlay not loaded"
            FAILED=$((FAILED+1))
        fi
    fi

    # Test 2: ALSA playback device
    echo -n "Test 2: ALSA playback... "
    if aplay -l 2>/dev/null | grep -qi "pcm1808\|pcm5102"; then
        info "OK (card 0 found)"
    else
        err "FAILED — no playback card"
        FAILED=$((FAILED+1))
    fi

    # Test 3: ALSA capture device
    echo -n "Test 3: ALSA capture... "
    if arecord -l 2>/dev/null | grep -qi "pcm1808\|hat"; then
        info "OK (capture device found)"
    else
        err "FAILED — no capture device"
        FAILED=$((FAILED+1))
    fi

    # Test 4: I2S pin muxing
    echo -n "Test 4: I2S GPIO mux... "
    PIN_OK=true
    for pin in 18 19 20 21; do
        if ! raspi-gpio get "$pin" 2>/dev/null | grep -q "ALT5"; then
            PIN_OK=false
        fi
    done
    if $PIN_OK; then
        info "OK (GPIO 18-21 in ALT5)"
    else
        warn "WARN — I2S pins not all in ALT5 (check pin conflicts)"
        warn "  GPIO 18-21 must be free for I2S"
    fi

    # Test 5: Quick playback test (sine wave, 1 second)
    echo -n "Test 5: Audio playback... "
    if speaker-test -t sine -f 440 -l 1 -r 48000 -c 2 &>/dev/null; then
        info "OK (sine wave played)"
    else
        err "FAILED — speaker-test error"
        FAILED=$((FAILED+1))
    fi

    # Test 6: Quick recording test
    echo -n "Test 6: Audio capture... "
    if arecord -d 2 -f cd -t wav /tmp/hat_test.wav 2>/dev/null; then
        FILESIZE=$(stat -c%s /tmp/hat_test.wav 2>/dev/null || echo 0)
        if [ "$FILESIZE" -gt 1000 ]; then
            info "OK (${FILESIZE} bytes recorded)"
        else
            warn "WARN — recording too small (${FILESIZE} bytes)"
        fi
    else
        err "FAILED — arecord error"
        FAILED=$((FAILED+1))
    fi

    echo ""
    if [ "$FAILED" -eq 0 ]; then
        info "All tests passed!"
    else
        warn "${FAILED} test(s) failed — check dmesg and wiring"
        exit 1
    fi
fi