#!/usr/bin/env bash
# ──────────────────────────────────────────────────────────────────────────────
# Pi Multi-FX Pedal — System Minimizer
#
# Profiles current resource usage, then strips away everything not needed
# for real-time guitar DSP. Run AFTER dietpi-postinstall.sh (or on any
# fresh DietPi/RPi OS install).
#
# This script is idempotent — safe to run multiple times. The uninstall
# list is documented below; review before running in production.
#
# Usage:
#   sudo ./minimize.sh [--dry-run] [--aggressive]
#
#   --dry-run     Only profile and list, don't remove anything
#   --aggressive  Also remove firmware, locales, docs, Python test packages
#
# Target: < 60 MB idle RAM, < 500 MB disk
# ──────────────────────────────────────────────────────────────────────────────

set -euo pipefail

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

for arg in "$@"; do
    case "$arg" in
        --dry-run)    DRY_RUN=true ;;
        --aggressive) AGGRESSIVE=true ;;
    esac
done

# ── 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 ./minimize.sh)"
    exit 1
fi

info "========== Pi Multi-FX Pedal — System Minimizer =========="
[[ "$DRY_RUN" == true ]] && warn "DRY RUN MODE — no changes will be made"
echo ""

# ══════════════════════════════════════════════════════════════════════════════
# Step 0: Baseline — record current state
# ══════════════════════════════════════════════════════════════════════════════

info "Step 0: Recording baseline..."

BASELINE_FILE="/tmp/pedal-minimize-baseline-$(date +%s).txt"
{
    echo "=== Baseline: $(date -Iseconds) ==="
    echo ""
    echo "--- Memory ---"
    free -m
    echo ""
    echo "--- Disk ---"
    df -h /
    echo ""
    echo "--- Top 30 packages by installed size ---"
    dpkg-query -Wf '${Installed-Size}\t${Package}\n' 2>/dev/null | sort -rn | head -30
    echo ""
    echo "--- Enabled systemd services ---"
    systemctl list-units --type=service --state=running --no-legend 2>/dev/null | awk '{print $1}' | sort
    echo ""
    echo "--- Boot time (systemd-analyze) ---"
    systemd-analyze 2>/dev/null || echo "not available"
    echo ""
    echo "--- Kernel modules loaded ---"
    lsmod | sort
} > "$BASELINE_FILE"

ok "Baseline saved to $BASELINE_FILE"

# ══════════════════════════════════════════════════════════════════════════════
# Step 1: Remove unnecessary packages
# ══════════════════════════════════════════════════════════════════════════════

info "Step 1: Removing unnecessary packages..."

# ── Packages we KNOW we don't need on a headless audio pedal ──────────────
PACKAGES_TO_PURGE=(
    # Bluetooth (no BT on a pedal)
    bluez bluez-tools bluez-firmware pi-bluetooth
    # Printing
    cups cups-bsd cups-client cups-common cups-core-drivers cups-daemon cups-filters cups-ppdc cups-server-common
    # Avahi/mDNS (not needed for local audio)
    avahi-daemon avahi-autoipd avahi-utils
    # ModemManager
    modemmanager
    # Triggerhappy (input daemon — we handle GPIO ourselves)
    triggerhappy
    # X11 / display server (headless — OLED is SPI/I2C)
    xserver-xorg xserver-xorg-core xserver-xorg-legacy x11-common
    # Desktop environments (shouldn't be installed on minimal DietPi, but just in case)
    raspberrypi-ui-mods raspberrypi-net-mods
    lxde lxde-common lxmenu-data lxterminal
    # Office/libreoffice
    libreoffice* liblibreoffice*
    # Web browsers
    chromium-browser chromium-browser-l10n chromium-codecs-ffmpeg
    firefox-esr
    # Games
    minecraft-pi python3-pgzero
    # Sound servers that conflict with JACK
    pulseaudio pulseaudio-utils pulseaudio-module-bluetooth
    pipewire pipewire-audio pipewire-pulse
    # WiFi firmware (if wired-only — we prompt below)
    # Removed conditionally if --aggressive
)

# Actually remove them
for pkg in "${PACKAGES_TO_PURGE[@]}"; do
    if dpkg -l "$pkg" &>/dev/null 2>&1; then
        if [[ "$DRY_RUN" == true ]]; then
            info "[DRY RUN] Would purge: $pkg"
        else
            apt-get purge -y -qq "$pkg" 2>/dev/null || true
            ok "Purged: $pkg"
        fi
    fi
done

# ── Aggressive mode: remove even more ──────────────────────────────────────
if [[ "$AGGRESSIVE" == true ]]; then
    info "Aggressive mode: removing additional packages..."

    AGGRESSIVE_PACKAGES=(
        # Firmware (keep only what's needed for boot + I2S)
        firmware-brcm80211 firmware-atheros firmware-iwlwifi firmware-realtek
        # Locales (keep only en_US.UTF-8)
        locales-all language-pack-* language-pack-en-base
        # Documentation
        man-db manpages info
        # Python test packages (not needed at runtime)
        python3-pytest python3-pytest-asyncio python3-pytest-mock
        # X11 fonts
        fonts-* xfonts-*
        # Unused sound modules for hardware we don't have
        alsa-base  # we use our own /etc/asound.conf
        # Unused Python packages on DietPi
        python3-pygame python3-tk
    )

    for pkg in "${AGGRESSIVE_PACKAGES[@]}"; do
        if dpkg -l "$pkg" &>/dev/null 2>&1; then
            if [[ "$DRY_RUN" == true ]]; then
                info "[DRY RUN] Would purge: $pkg"
            else
                apt-get purge -y -qq "$pkg" 2>/dev/null || true
                ok "Purged: $pkg"
            fi
        fi
    done

    # Remove unused locale data
    if [[ "$DRY_RUN" == false ]]; then
        localedef --list-archive 2>/dev/null | grep -v en_US | while read -r locale; do
            localedef --delete-from-archive "$locale" 2>/dev/null || true
        done
        ok "Stripped non-English locales"
    fi
fi

# Clean up
if [[ "$DRY_RUN" == false ]]; then
    apt-get autoremove --purge -y -qq 2>/dev/null || true
    apt-get autoclean -y -qq 2>/dev/null || true
    ok "apt autoremove/autoclean complete"
fi

# ══════════════════════════════════════════════════════════════════════════════
# Step 2: Disable unnecessary systemd services
# ══════════════════════════════════════════════════════════════════════════════

info "Step 2: Disabling unnecessary systemd services..."

SERVICES_TO_DISABLE=(
    bluetooth
    avahi-daemon
    triggerhappy
    ModemManager
    systemd-resolved
    apt-daily.timer
    apt-daily-upgrade.timer
    man-db.timer
    e2scrub_all.timer
    fstrim.timer
)

for svc in "${SERVICES_TO_DISABLE[@]}"; do
    if systemctl is-enabled "$svc" &>/dev/null 2>&1; then
        if [[ "$DRY_RUN" == true ]]; then
            info "[DRY RUN] Would disable: $svc"
        else
            systemctl disable --now "$svc" 2>/dev/null || true
            ok "Disabled: $svc"
        fi
    fi
done

# Mask pulseaudio completely (not just disable)
for unit in pulseaudio.service pulseaudio.socket; do
    if systemctl list-unit-files "$unit" &>/dev/null 2>&1; then
        if [[ "$DRY_RUN" == true ]]; then
            info "[DRY RUN] Would mask: $unit"
        else
            systemctl mask "$unit" 2>/dev/null || true
            ok "Masked: $unit"
        fi
    fi
done

# ══════════════════════════════════════════════════════════════════════════════
# Step 3: Remove unused kernel modules
# ══════════════════════════════════════════════════════════════════════════════

info "Step 3: Removing unnecessary kernel modules..."

BLACKLIST_FILE="/etc/modprobe.d/pedal-blacklist.conf"
if [[ "$DRY_RUN" == false ]]; then
    cat > "$BLACKLIST_FILE" <<'BLACKLIST'
# Pi Multi-FX Pedal — blacklist unused kernel modules
# Audio subsystems we don't use
blacklist snd_bcm2835
blacklist snd_usb_audio
blacklist snd_usbmidi_lib
blacklist snd_rawmidi
blacklist snd_seq_device
blacklist snd_hwdep
blacklist snd_hda_intel
blacklist snd_hda_codec
blacklist snd_hda_core

# Bluetooth
blacklist bluetooth
blacklist btusb
blacklist btbcm

# WiFi (wired-only mode)
blacklist brcmfmac
blacklist brcmutil

# Unused input devices
blacklist joydev
blacklist uvcvideo
blacklist videobuf2_v4l2

# Unused storage
blacklist r8152   # Realtek USB NIC
blacklist xhci_pci # USB 3.0 (not needed on headless pedal)
BLACKLIST
    ok "Kernel module blacklist: $BLACKLIST_FILE"
else
    info "[DRY RUN] Would write: $BLACKLIST_FILE"
fi

# ══════════════════════════════════════════════════════════════════════════════
# Step 4: Profile and measure
# ══════════════════════════════════════════════════════════════════════════════

info "Step 4: Profiling post-minimization footprint..."

AFTER_FILE="/tmp/pedal-minimize-after-$(date +%s).txt"
{
    echo "=== After Minimize: $(date -Iseconds) ==="
    echo ""
    echo "--- Memory ---"
    free -m
    echo ""
    echo "--- Disk ---"
    df -h /
    echo ""
    echo "--- Top 10 packages by installed size ---"
    dpkg-query -Wf '${Installed-Size}\t${Package}\n' 2>/dev/null | sort -rn | head -10
    echo ""
    echo "--- Running services ---"
    systemctl list-units --type=service --state=running --no-legend 2>/dev/null | awk '{print $1}' | sort
    echo ""
    echo "--- systemd-analyze blame ---"
    systemd-analyze blame 2>/dev/null | head -20 || echo "not available"
} > "$AFTER_FILE"

ok "Post-minimize profile saved to $AFTER_FILE"

# ── Compute and display deltas ────────────────────────────────────────────────
echo ""
info "========== Before / After Summary =========="

# RAM
BEFORE_MEM=$(grep "^Mem:" "$BASELINE_FILE" 2>/dev/null | awk '{print $4}')
AFTER_MEM=$(grep "^Mem:" "$AFTER_FILE" 2>/dev/null | awk '{print $4}')
BEFORE_MEM_TOTAL=$(grep "^Mem:" "$BASELINE_FILE" 2>/dev/null | awk '{print $2}')
echo "  RAM:   ${BEFORE_MEM:-?} → ${AFTER_MEM:-?} MB free (of ${BEFORE_MEM_TOTAL:-?} MB total)"

# Disk
BEFORE_DISK=$(grep "^/dev/root" "$BASELINE_FILE" 2>/dev/null | awk '{print $4}')
AFTER_DISK=$(grep "^/dev/root" "$AFTER_FILE" 2>/dev/null | awk '{print $4}')
echo "  Disk:  ${BEFORE_DISK:-?} → ${AFTER_DISK:-?} available"

# Package count
BEFORE_PKGS=$(dpkg-query -f '.\n' -W 2>/dev/null | wc -l)
echo "  Packages: $BEFORE_PKGS installed (after removal)"

echo ""
if [[ "$DRY_RUN" == true ]]; then
    warn "DRY RUN — no changes were made"
    info "Re-run without --dry-run to apply changes"
else
    ok "Minimization complete!"
    info "Baseline: $BASELINE_FILE"
    info "After:    $AFTER_FILE"
    echo ""
    info "A reboot is recommended to fully apply kernel module blacklist:"
    echo "  sudo reboot"
    echo ""
    info "After reboot, verify footprint:"
    echo "  free -m | grep Mem"
    echo "  df -h /"
    echo "  systemd-analyze"
fi