#!/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_irq_affinity() {
    local irq="$1"
    local label="$2"
    local cpu_list="$3"
    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 ($label) 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 ($label)" >&2
        return 1
    }

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

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

    irqs="$(grep -i "$device" /proc/interrupts | awk '{print $1}' | sed 's/://')" || true

    if [[ -z "$irqs" ]]; then
        echo "Warning: IRQ for $device not found, skipping." >&2
        return 0
    fi

    local irq
    while IFS= read -r irq; do
        set_irq_affinity "$irq" "$device" "$cpu_list" || return 1
    done <<< "$irqs"
}

pibridge_uart_node() {
    # Device-tree uart node with the pibridge serdev child, empty on devices
    # without a pibridge (e.g. RevPi Flat S).
    local compat
    compat="$(grep -rl 'kunbus,pi-bridge' /proc/device-tree 2>/dev/null | grep '/compatible$' | head -1)"
    [[ -n "$compat" ]] || return 1
    dirname "$(dirname "$compat")"
}

pibridge_pl011_irq() {
    # On CM5 each pl011 has its own IRQ, all named "uart-pl011". Match the
    # pibridge uart's device-tree hwirq to pick the right one. Empty if it
    # can't be matched, e.g. the single shared IRQ on CM4.
    local node="$1"
    local intfile="$node/interrupts"
    local -a bytes cand
    [[ -r "$intfile" ]] || return 1

    # The hwirq is among the interrupt specifier cells.
    read -r -a bytes < <(od -An -v -tx1 "$intfile" | tr '\n' ' ')
    local i word
    for ((i = 0; i + 3 < ${#bytes[@]}; i += 4)); do
        word="${bytes[i]}${bytes[i+1]}${bytes[i+2]}${bytes[i+3]}"
        cand+=( "$((16#$word))" )
    done

    local d hw v
    for d in /sys/kernel/irq/*; do
        grep -q 'uart-pl011' "$d/actions" 2>/dev/null || continue
        hw="$(cat "$d/hwirq" 2>/dev/null)"
        for v in "${cand[@]}"; do
            if [[ "$hw" == "$v" ]]; then
                basename "$d"
                return 0
            fi
        done
    done

    return 1
}

set_picontrol_affinity() {
    # Co-locate the pibridge pl011 uart IRQ and the piControl I/O task on the
    # same cpu. Devices without a pibridge (e.g. RevPi Flat S) have nothing to
    # pin. Resolve the specific IRQ on CM5, otherwise fall back to the
    # "uart-pl011" name match for the single shared IRQ on CM4.
    local cpu_list="$1"
    local node irq

    node="$(pibridge_uart_node)" || return 0

    irq="$(pibridge_pl011_irq "$node")" || true
    if [[ -n "$irq" ]]; then
        set_irq_affinity "$irq" "uart-pl011 (pibridge)" "$cpu_list"
    else
        set_affinity "uart-pl011" "$cpu_list"
    fi

    taskset -cp "$cpu_list" $(pgrep piControl\ I\/O)
}

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

set_picontrol_affinity "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"
