#!/usr/bin/env bash
# bc250-cu-mask.sh - generate/install selective WGP masks for BC-250 40-CU unlock.

set -euo pipefail

RESULTS="${BC250_CU_HEALTH_RESULTS:-/var/lib/bc250-cu-health-test/results.tsv}"
CONF="${BC250_CU_MASK_CONF:-/etc/modprobe.d/bc250-40cu-selective-mask.conf}"
INSTALL=0
BAD_ITEMS=()
TMP_CONF="$(mktemp)"
trap 'rm -f "$TMP_CONF"' EXIT

usage() {
	cat <<EOF
Usage: $0 [--results FILE] [--bad SE.SH.WGP[,SE.SH.WGP...]] [--bad-cu SE.SH.CU[,SE.SH.CU...]] [--install]

Examples:
  $0 --bad 1.0.3
  $0 --results /var/lib/bc250-cu-health-test/results.tsv --install
  $0 --bad-cu 1.0.6,1.0.7 --install

Output modprobe config:
  options amdgpu bc250_cc_write_mode=3 disable_cu=SE.SH.CU,...

WGP is the hardware disable granularity. Each bad WGP expands to both CUs.
EOF
}

die() {
	echo "ERROR: $*" >&2
	exit 1
}

split_csv() {
	local item
	IFS=',' read -ra _items <<<"$1"
	for item in "${_items[@]}"; do
		[ -n "$item" ] && BAD_ITEMS+=("$item")
	done
}

while [ "$#" -gt 0 ]; do
	case "$1" in
		--results)
			RESULTS="${2:?missing value for --results}"
			shift 2
			;;
		--bad)
			split_csv "${2:?missing value for --bad}"
			shift 2
			;;
		--bad-cu)
			IFS=',' read -ra _cus <<<"${2:?missing value for --bad-cu}"
			for cu_item in "${_cus[@]}"; do
				IFS='.' read -r se sh cu <<<"$cu_item"
				[[ "$se" =~ ^[0-1]$ && "$sh" =~ ^[0-1]$ && "$cu" =~ ^[0-9]$ ]] ||
					die "invalid CU entry: $cu_item"
				BAD_ITEMS+=("$se.$sh.$((cu / 2))")
			done
			shift 2
			;;
		--install)
			INSTALL=1
			shift
			;;
		-h|--help)
			usage
			exit 0
			;;
		*)
			die "unknown argument: $1"
			;;
	esac
done

if [ "${#BAD_ITEMS[@]}" -eq 0 ] && [ -f "$RESULTS" ]; then
	while IFS=$'\t' read -r idx se sh wgp status rest; do
		case "$idx" in ""|\#*) continue ;; esac
		[ "$status" = "FAIL" ] || continue
		BAD_ITEMS+=("$se.$sh.$wgp")
	done <"$RESULTS"
fi

declare -A seen_wgp=()
declare -A seen_cu=()
bad_wgps=()
disable_cus=()

for item in "${BAD_ITEMS[@]}"; do
	IFS='.' read -r se sh wgp <<<"$item"
	[[ "$se" =~ ^[0-1]$ && "$sh" =~ ^[0-1]$ && "$wgp" =~ ^[0-4]$ ]] ||
		die "invalid WGP entry: $item"
	key="$se.$sh.$wgp"
	if [ -z "${seen_wgp[$key]+x}" ]; then
		seen_wgp[$key]=1
		bad_wgps+=("$key")
	fi
	for cu in $((wgp * 2)) $((wgp * 2 + 1)); do
		cu_key="$se.$sh.$cu"
		if [ -z "${seen_cu[$cu_key]+x}" ]; then
			seen_cu[$cu_key]=1
			disable_cus+=("$cu_key")
		fi
	done
done

disable_csv="$(IFS=,; echo "${disable_cus[*]}")"

{
	echo "# BC-250 selective 40-CU mask."
	echo "# Generated by $0 on $(date -Iseconds)."
	echo "# Bad WGPs: ${bad_wgps[*]:-none}"
	echo "# Requires patched amdgpu with bc250_cc_write_mode support."
	if [ -n "$disable_csv" ]; then
		echo "options amdgpu bc250_cc_write_mode=3 disable_cu=$disable_csv"
	else
		echo "options amdgpu bc250_cc_write_mode=3"
	fi
} >"$TMP_CONF"

if [ "$INSTALL" -eq 1 ]; then
	[ "$(id -u)" = "0" ] || die "--install requires root"
	install -m 0644 "$TMP_CONF" "$CONF"
	echo "Installed $CONF"
	if command -v update-initramfs >/dev/null 2>&1; then
		update-initramfs -u -k "$(uname -r)" || true
	elif command -v dracut >/dev/null 2>&1; then
		dracut -f || true
	fi
else
	cat "$TMP_CONF"
fi

usable=$((40 - ${#bad_wgps[@]} * 2))
echo "# Usable after mask: $usable/40 CUs (${#bad_wgps[@]} bad WGPs masked)"