#!/bin/bash
#
# SPDX-License-Identifier: GPL-2.0-or-later
#
# SPDX-FileCopyrightText: 2025 KUNBUS GmbH

set -euo pipefail

if [[ $EUID -ne 0 ]]; then
    echo "Error: This script must be run as root" >&2
    exit 1
fi

unsupported_cm() {
    # Only Compute Modules with BCM2711 or newer support setting the IRQ
    # CPU affinity. This function checks against known unsupported modules.
    #
    # List of revision codes:
    # https://www.raspberrypi.com/documentation/computers/raspberry-pi.html#new-style-revision-codes

    local revision=$(awk '/^Revision/ {print $3}' /proc/cpuinfo)

    case "$revision" in
        900061 | a020a0 | a220a0 | a02100)
            # CM1, CM3, CM3+
            return 0 ;;
        *)
            # Other compute modules (e.g. CM4, CM4S, CM5 or newer)
            return 1 ;;
    esac
}

set_affinity() {
    local device="$1"
    local cpu_list="$2"
    local irq

    irq="$(grep -i "$device" /proc/interrupts | awk '{print $1}' | sed 's/://')" || {
        echo "Warning: IRQ for $device not found, skipping." >&2
        return 0
    }

    local current_affinity
    current_affinity="$(cat /proc/irq/"$irq"/smp_affinity_list 2>/dev/null || echo "")"

    if [[ "$current_affinity" == "$cpu_list" ]]; then
        echo "IRQ $irq ($device) is already set to CPU(s) $cpu_list, skipping."
        return 0
    fi

    echo "$cpu_list" > /proc/irq/"$irq"/smp_affinity_list || {
        echo "Error: Failed to set affinity for IRQ $irq ($device)" >&2
        return 1
    }

    echo "Set IRQ $irq ($device) to CPU(s) $cpu_list"
}

if unsupported_cm; then
    echo "CM1, CM3, CM3+ don't support setting CPU affinity. Exiting."
    exit 0
fi

# Run IRQ handler for pl011 uart on cpu core 3
set_affinity "uart-pl011" "3"

# Run IRQ handler for mmc* (eMMC and if present wifi) on cpu core 2
set_affinity "mmc0" "2"
# mmc1 (WiFi) is attached to a separate IRQ on CM5
set_affinity "mmc1" "2"
