#!/bin/bash
# Apply CPU cores defaults to cmdline.txt (fallback for manual installations)
# Used by dangerous-pi-cpu-cores.service
#
# NOTE: When using the official pi-gen image, maxcpus is already set during build.
# This script only runs as a fallback for manual installations.
#
# On Pi Zero 2 W (and other ARM devices), CPU hotplug via sysfs does NOT work.
# Instead, we use the maxcpus=N kernel parameter in cmdline.txt.
# This script sets the default maxcpus for Pi Zero 2 W if not already configured.

LOG_TAG="dangerous-pi-cpu-cores"

log() {
    logger -t "$LOG_TAG" "$1"
    echo "$1"
}

# Find cmdline.txt (location varies by Pi OS version)
find_cmdline() {
    if [ -f /boot/firmware/cmdline.txt ]; then
        echo "/boot/firmware/cmdline.txt"
    elif [ -f /boot/cmdline.txt ]; then
        echo "/boot/cmdline.txt"
    else
        echo ""
    fi
}

CMDLINE_FILE=$(find_cmdline)

if [ -z "$CMDLINE_FILE" ]; then
    log "ERROR: Could not find cmdline.txt"
    exit 1
fi

# Check if maxcpus is already configured
if grep -q "maxcpus=" "$CMDLINE_FILE"; then
    CURRENT=$(grep -o 'maxcpus=[0-9]*' "$CMDLINE_FILE" | cut -d= -f2)
    log "maxcpus already configured: $CURRENT cores"
    exit 0
fi

# Determine default cores based on Pi model
get_default_cores() {
    local model=""
    local total_cores=$(nproc --all)

    if [ -f /proc/device-tree/model ]; then
        model=$(cat /proc/device-tree/model | tr -d '\0')
    fi

    # Pi Zero 2 W has 4 cores but should default to 2 for thermal management
    if echo "$model" | grep -q "Pi Zero 2"; then
        echo 2
    # Single-core Pis - no need to set maxcpus
    elif echo "$model" | grep -q "Pi Zero W\|Pi Zero$\|Pi 1\|Pi Model A\|Pi Model B"; then
        echo "$total_cores"  # Return total (no restriction needed)
    # All other multi-core Pis use all cores by default
    else
        echo "$total_cores"
    fi
}

DEFAULT_CORES=$(get_default_cores)
TOTAL_CORES=$(nproc --all)

# Only set maxcpus if we want fewer than total cores
if [ "$DEFAULT_CORES" -lt "$TOTAL_CORES" ]; then
    log "Setting default maxcpus=$DEFAULT_CORES for this device"

    # Backup original
    cp "$CMDLINE_FILE" "${CMDLINE_FILE}.bak"

    # Append maxcpus to cmdline
    CURRENT_CMDLINE=$(cat "$CMDLINE_FILE" | tr -d '\n')
    echo "$CURRENT_CMDLINE maxcpus=$DEFAULT_CORES" > "${CMDLINE_FILE}.new"
    mv "${CMDLINE_FILE}.new" "$CMDLINE_FILE"

    log "Updated $CMDLINE_FILE with maxcpus=$DEFAULT_CORES"
    log "NOTE: This will take effect on next reboot"
else
    log "No maxcpus restriction needed (using all $TOTAL_CORES cores)"
fi

exit 0