#!/usr/bin/env bash
# ────────────────────────────────────────────────────────────────────
# Pi Multi-FX Pedal — RT Performance Tuning Script
#
# Applies real-time audio optimizations for stable guitar playing on
# RPi 4B. Designed to run at boot (via systemd or main.py) and on
# every JACK/audio profile change.
#
# Targets:
#   - Round-trip latency: <12ms (ideally <8ms)
#   - Zero xruns during playing
#   - CPU <40% on Pi 4B at recommended settings (512/48k)
#
# Usage:
#   sudo ./rt-tune.sh                    # apply all optimizations
#   sudo ./rt-tune.sh --status           # check current settings
#   sudo ./rt-tune.sh --irq-only         # just IRQ affinity
#   sudo ./rt-tune.sh --usb-audio-irq    # find + pin USB audio IRQ
# ────────────────────────────────────────────────────────────────────

set -euo pipefail

# ── 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" "$*"; }

# ── Config ─────────────────────────────────────────────────────────
# Which CPU core to dedicate to audio IRQ handling
# RPi 4B has 4 cores (0-3). Core 3 is largely unused by Linux housekeeping.
# Core 0 = most interrupts + kernel, Core 1-2 = general, Core 3 = isolated
IRQ_CORE=3

# Default JACK parameters (will be overridden by main.py dynamically)
JACK_PERIOD=512
JACK_RATE=48000

# ── Helpers ─────────────────────────────────────────────────────────

# Find the USB audio interface IRQ number
find_usb_audio_irq() {
    # Look for the USB audio device in /proc/interrupts
    # On RPi 4B, the USB controller is on a PCIe bridge (xhci-hcd) or
    # directly on the BCM2711's DWC2/dwc_otg USB controller.
    # Focusrite Scarlett 2i2 shows up as a USB interrupt tied to xhci-hcd.

    # Strategy 1: Look for xhci-hcd (USB 3.0 controller on Pi 4B)
    local irq
    irq=$(awk '/xhci-hcd/ {gsub(":","",$1); print $1}' /proc/interrupts 2>/dev/null | head -1)

    # Strategy 2: Look for dwc_otg/dwc2 (USB 2.0 controller)
    if [[ -z "$irq" ]]; then
        irq=$(awk '/dwc_otg|dwc2/ {gsub(":","",$1); print $1}' /proc/interrupts 2>/dev/null | head -1)
    fi

    # Strategy 3: Find the USB controller from sysfs for the audio device
    if [[ -z "$irq" ]]; then
        # Try to find the USB device that's our audio interface
        # Look for USB audio class devices
        local usb_dev
        usb_dev=$(grep -l "audio" /sys/bus/usb/devices/*/bInterfaceClass 2>/dev/null | head -1)
        if [[ -n "$usb_dev" ]]; then
            local usb_path
            usb_path=$(dirname "$usb_dev")
            # Walk up to find the parent USB controller
            while [[ "$usb_path" != "/sys/bus/usb/devices" && "$usb_path" != "/" ]]; do
                if [[ -f "$usb_path/irq" ]]; then
                    irq=$(cat "$usb_path/irq" 2>/dev/null)
                    break
                fi
                usb_path=$(dirname "$usb_path" 2>/dev/null)
            done
        fi
    fi

    echo "$irq"
}

# ── Apply IRQ affinity ──────────────────────────────────────────────

apply_irq_affinity() {
    local irq="$1"
    local core="$2"
    local smp_affinity

    if [[ -z "$irq" || "$irq" == "0" ]]; then
        warn "No USB audio IRQ found — cannot set affinity"
        return 1
    fi

    # Convert core number to hex mask for /proc/irq/*/smp_affinity
    # Core 0 = 1, Core 1 = 2, Core 2 = 4, Core 3 = 8
    smp_affinity=$(printf "%x" $((1 << core)))

    local irq_dir="/proc/irq/$irq"
    if [[ ! -d "$irq_dir" ]]; then
        warn "IRQ $irq directory not found at $irq_dir"
        return 1
    fi

    echo "$smp_affinity" > "$irq_dir/smp_affinity" 2>/dev/null || true
    local effective
    effective=$(cat "$irq_dir/smp_affinity" 2>/dev/null || echo "unknown")

    # Also set smp_affinity_list for convenience
    echo "$core" > "$irq_dir/smp_affinity_list" 2>/dev/null || true

    info "IRQ $irq pinned to core $core (smp_affinity=0x$smp_affinity, effective=0x$effective)"
    return 0
}

# ── Set RT priority for the current process ──────────────────────────

set_self_rt_priority() {
    local prio="$1"
    # Set SCHED_FIFO with the given priority for our process group
    # Note: This only works if called from the target process or with CAP_SYS_NICE
    chrt -f -p "$prio" $$ 2>/dev/null || warn "Cannot set self RT priority (not running as root?)"
}

# ── CPU isolation ───────────────────────────────────────────────────

# Move all non-critical interrupts away from the IRQ core
isolate_core_from_housekeeping() {
    local core="$1"
    local exclude_irqs="$2"  # comma-separated IRQs to keep on the isolated core

    # For each IRQ, move it away from our dedicated audio core
    # unless it's one we explicitly want to keep there
    for irq_dir in /proc/irq/[0-9]*/; do
        local irq_num
        irq_num=$(basename "$irq_dir")
        local irq_name
        irq_name=$(cat "${irq_dir}affinity_hint" 2>/dev/null || echo "")

        # Skip if this is the USB audio IRQ we want to keep pinned
        if [[ ",$exclude_irqs," == *",$irq_num,"* ]]; then
            continue
        fi

        # Move to other cores (0-2, leaving core 3 free)
        # Affinity = 0x7 (cores 0,1,2)
        echo "7" > "${irq_dir}smp_affinity" 2>/dev/null || true
    done
    info "Non-audio IRQs moved away from core $core"
}

# ── Status check ────────────────────────────────────────────────────

show_status() {
    echo ""
    info "========== RT Performance Status =========="
    echo ""

    # CPU governor
    echo "── CPU Governor ──"
    for c in /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor; do
        local cpu
        cpu=$(basename "$(dirname "$c")")
        echo "  $cpu: $(cat "$c" 2>/dev/null || echo 'N/A')"
    done

    # Current frequency
    echo ""
    echo "── CPU Frequency ──"
    for c in /sys/devices/system/cpu/cpu*/cpufreq/scaling_cur_freq; do
        local cpu
        cpu=$(basename "$(dirname "$c")")
        local freq_khz
        freq_khz=$(cat "$c" 2>/dev/null || echo 'N/A')
        echo "  $cpu: $((freq_khz / 1000)) MHz"
    done

    # IRQ affinity
    echo ""
    echo "── IRQ Affinity ──"
    local irq
    irq=$(find_usb_audio_irq)
    if [[ -n "$irq" && "$irq" != "0" ]]; then
        local aff
        aff=$(cat "/proc/irq/$irq/smp_affinity" 2>/dev/null || echo "N/A")
        local name
        name=$(cat "/proc/irq/$irq/name" 2>/dev/null || echo "unknown")
        echo "  USB Audio IRQ $irq ($name): smp_affinity=0x$aff"
    else
        echo "  No USB audio IRQ found"
    fi

    # All IRQs
    echo ""
    echo "── All Interrupts (affinity) ──"
    for irq_dir in /proc/irq/[0-9]*/; do
        local i
        i=$(basename "$irq_dir")
        local aff
        aff=$(cat "${irq_dir}smp_affinity" 2>/dev/null || echo "?")
        local name
        name=$(cat "${irq_dir}name" 2>/dev/null || echo "?")
        printf "  IRQ %-4s  0x%-4s  %s\\n" "$i" "$aff" "$name"
    done

    # JACK status
    echo ""
    echo "── JACK Status ──"
    if pidof jackd >/dev/null 2>&1; then
        echo "  jackd: RUNNING"
        # Show JACK command line
        ps aux | grep jackd | grep -v grep | head -1 | awk '{$1=$2=$3=$4=$5=$6=$7=$8=$9=$10=""; print "  Args:" $0}'
    else
        echo "  jackd: NOT RUNNING"
    fi

    # RT priority of processes
    echo ""
    echo "── RT Priority ──"
    for proc in jackd python3; do
        pidof "$proc" 2>/dev/null | tr ' ' '\\n' | while read -r pid; do
            local policy
            policy=$(chrt -p "$pid" 2>/dev/null | head -1 || echo "N/A")
            echo "  $proc (PID $pid): $policy"
        done
    done

    # Memory locking
    echo ""
    echo "── Memory Locking ──"
    if [[ -f /proc/self/status ]]; then
        grep -i "lock" /proc/self/status 2>/dev/null | head -5 | while read -r line; do
            echo "  $line"
        done
    fi

    # XRun debug
    echo ""
    echo "── XRun Debug ──"
    for card in /proc/asound/card*/xrun_debug; do
        echo "  $card: $(cat "$card" 2>/dev/null || echo 'N/A')"
    done

    echo ""
    ok "Status check complete"
}

# ═══════════════════════════════════════════════════════════════════
# Main
# ═══════════════════════════════════════════════════════════════════

if [[ $EUID -ne 0 ]]; then
    err "This script must be run as root"
    exit 1
fi

MODE="${1:-all}"

case "$MODE" in
    --status|-s)
        show_status
        exit 0
        ;;

    --irq-only|-i)
        IRQ=$(find_usb_audio_irq)
        if [[ -n "$IRQ" && "$IRQ" != "0" ]]; then
            apply_irq_affinity "$IRQ" "$IRQ_CORE" && ok "IRQ affinity applied"
        else
            warn "No USB audio IRQ found — checking /proc/interrupts..."
            grep -E "xhci|dwc|usb|audio" /proc/interrupts | head -10
            echo ""
            info "To find the right IRQ manually:"
            info "  cat /proc/interrupts | grep xhci"
            info "  cat /sys/bus/usb/devices/*/irq 2>/dev/null"
        fi
        exit 0
        ;;

    --usb-audio-irq|-u)
        IRQ=$(find_usb_audio_irq)
        if [[ -n "$IRQ" && "$IRQ" != "0" ]]; then
            local name
            name=$(cat "/proc/irq/$IRQ/name" 2>/dev/null || echo "unknown")
            info "USB Audio IRQ = $IRQ ($name)"
        else
            warn "No USB audio IRQ found"
            echo "── /proc/interrupts (USB/audio lines) ──"
            grep -E "xhci|dwc|usb|audio|snd" /proc/interrupts | head -10
        fi
        exit 0
        ;;

    all|*)
        # Full tuning
        info "========== Pi Multi-FX Pedal — RT Tuning =========="
        echo ""

        # ── 1. Find and pin USB audio IRQ ────────────────────────
        info "Step 1: USB audio IRQ affinity..."
        IRQ=$(find_usb_audio_irq)
        if [[ -n "$IRQ" && "$IRQ" != "0" ]]; then
            apply_irq_affinity "$IRQ" "$IRQ_CORE"
            # Isolate core 3 from housekeeping interrupts
            isolate_core_from_housekeeping "$IRQ_CORE" "$IRQ"
        else
            warn "USB audio IRQ not found — audio may work but without dedicated IRQ core"
        fi

        # ── 2. CPU governor (already done in main.py, but belt-and-braces) ──
        info "Step 2: CPU governor → performance..."
        echo performance | tee /sys/devices/system/cpu/cpu*/cpufreq/scaling_governor >/dev/null 2>&1 || \
            warn "Could not set CPU governor (may need cpufrequtils)"

        # ── 3. Disable CPU idle states (C-states) for lower latency ──
        # This prevents the CPU from entering deep sleep that adds latency
        # on wake. Trade-off: ~0.5W extra power draw.
        info "Step 3: Disabling deep CPU idle states..."
        if [[ -f /sys/module/processor/parameters/max_cstate ]]; then
            echo 1 > /sys/module/processor/parameters/max_cstate 2>/dev/null || true
        fi
        if [[ -f /dev/cpu_dma_latency ]]; then
            # Write 0 to request minimum DMA latency (blocks deep C-states)
            # The file stays open while the process lives — we do it briefly
            echo 0 > /dev/cpu_dma_latency 2>/dev/null || true
        fi

        # ── 4. Set JACK's ALSA buffer sizes for low latency ───────────
        # These are advisory — main.py overrides them via jackd arguments
        info "Step 4: ALSA buffer constraints (advisory)..."
        cat > /etc/security/limits.d/99-audio.conf <<'LIMITS'
# Pi Multi-FX Pedal — Real-time audio limits
# Applied by rt-tune.sh
@audio   -  rtprio     95
@audio   -  memlock    unlimited
@audio   -  nice      -20
LIMITS
        ok "Audio limits written to /etc/security/limits.d/99-audio.conf"

        # ── 5. Set xrun_debug for diagnostics ─────────────────────────
        info "Step 5: Enable xrun tracking..."
        for card in /proc/asound/card*/xrun_debug; do
            if [[ -w "$card" ]]; then
                # Bit 0 = enable xrun logging
                # Bit 1 = show stack backtrace on xrun
                # Bit 2 = inhibit xrun (test mode)
                # Value 3 = log xruns with backtrace (diagnostic)
                echo 3 > "$card" 2>/dev/null || true
                ok "Set xrun_debug on $(dirname "$card" | xargs basename)"
            fi
        done

        # ── 6. Process RT priority (for the calling process) ──────────
        info "Step 6: Setting RT priority..."
        # This is primarily done by main.py (mlockall, GC disable)
        # and the systemd service (LimitRTPRIO). We set it here too
        # so the script works in all contexts.
        set_self_rt_priority 80 || true

        # ── 7. Systemd service check ──────────────────────────────────
        info "Step 7: Verifying systemd RT limits..."
        if systemctl is-active pi-multifx-pedal.service &>/dev/null; then
            local rtprio
            rtprio=$(systemctl show pi-multifx-pedal.service -p LimitRTPRIO --value 2>/dev/null || echo "?")
            local memlock
            memlock=$(systemctl show pi-multifx-pedal.service -p LimitMEMLOCK --value 2>/dev/null || echo "?")
            local nice
            nice=$(systemctl show pi-multifx-pedal.service -p LimitNICE --value 2>/dev/null || echo "?")
            info "  LimitRTPRIO=$rtprio  LimitMEMLOCK=$memlock  LimitNICE=$nice"
            if [[ "$rtprio" != "95" ]]; then
                warn "LimitRTPRIO should be 95 — check pi-multifx-pedal.service"
            fi
            if [[ "$memlock" != "infinity" ]]; then
                warn "LimitMEMLOCK should be infinity — check pi-multifx-pedal.service"
            fi
        else
            warn "pi-multifx-pedal.service not running — check after deployment"
        fi

        echo ""
        ok "========== RT Tuning Complete =========="
        echo ""
        info "Recommended next steps:"
        info "  1. Start pedal:  sudo systemctl start pi-multifx-pedal.service"
        info "  2. Check logs:   journalctl -fu pi-multifx-pedal.service"
        info "  3. Test latency: jack_iodelay (needs loopback cable)"
        info "  4. Check xruns:  cat /proc/asound/card*/xrun_debug"
        info "  5. Full status:  sudo ./rt-tune.sh --status"
        ;;
esac