#!/bin/bash

watt_limit=120

for i in {0..10}; do
    card="/sys/class/drm/card$i"
    # Skip non-existing, virtual or non-GPU devices
    if [ ! -d "$card" ] || [ ! -e "$card/device" ]; then
        continue
    fi

    # Resolve the PCI device path
    device_path=$(readlink -f "$card/device")

    # Read vendor and device IDs from sysfs
    vendor_id=$(cat "$device_path/vendor")
    device_id=$(cat "$device_path/device")

    if [ "$vendor_id" == "0x1002" ] && [ "$device_id" == "0x66a0" ]; then 
	    echo "Vega detected, limiting to $watt_limit W."

	    # TODO check if the hwmon / power indices are stable with multiple cards or 
	    # make the script robust to always set the correct power cap
	    found=0
	    for hwmon in /sys/class/hwmon/hwmon*; do
        	if [ -e "$hwmon/device" ]; then
                	hwmon_dev=$(readlink -f "$hwmon/device")
                	if [ "$hwmon_dev" = "$device_path" ]; then
                    		power_cap_file="$hwmon/power1_cap"
                    		if [ -e "$power_cap_file" ]; then
                        		echo ${watt_limit}000000 > $power_cap_file
                        		found=1
                        		break
                    		fi
                	fi
            	fi
	    done

            if [ $found -eq 0 ]; then
               echo "Warning: Could not find power1_cap for this GPU."
            fi
    fi 

    echo
done