#!/usr/bin/env bash

IMAGE_INFO="/usr/share/ublue-os/image-info.json"
IMAGE_NAME=$(jq -r '."image-name"' < $IMAGE_INFO)
BASE_IMAGE_NAME=$(jq -r '."base-image-name"' < $IMAGE_INFO)

# SCRIPT VERSION
HWS_VER=1
HWS_VER_FILE="/etc/bazzite/hws_version"
HWS_VER_RAN=$(cat $HWS_VERSION_FILE)

# IMAGE IDENTIFIERS
KNOWN_IMAGE_NAME_FILE="/etc/bazzite/image_name"
KNOWN_IMAGE_NAME=$(cat $KNOWN_IMAGE_NAME_FILE)
KNOWN_IMAGE_FLAVOR_FILE="/etc/bazzite/image_flavor"
KNOWN_IMAGE_FLAVOR=$(cat $KNOWN_IMAGE_FLAVOR_FILE)

# Run script if updated
if [[ -f $HWS_VER_FILE && $HWS_VER = $HWS_VER_RAN ]]; then
  if [[ -f $KNOWN_IMAGE_NAME_FILE && -f $KNOWN_IMAGE_FLAVOR_FILE ]]; then
    # Run script if image has been rebased
    if [[ $IMAGE_NAME = $KNOWN_IMAGE_NAME && $IMAGE_FLAVOR = $KNOWN_IMAGE_FLAVOR ]]; then
      echo "Hardware setup has already ran. Exiting..."
      exit 0
    fi
  fi
fi

# GLOBAL
SYS_ID="$(cat /sys/devices/virtual/dmi/id/product_name)"
GPU_ID=$(lspci -k | grep -A 3 -E "(VGA|3D)")
KARGS=$(rpm-ostree kargs)
NEEDED_KARGS=""
echo "Current kargs: $KARGS"
mkdir -p /etc/bazzite

# FSTAB CONFIGURATION
if [[ $(grep "compress=zstd" /etc/fstab) ]]; then
  echo "Applying fstab param adjustments"
  if grep -q '64GB' <<< $(lsblk -o MODEL); then
    echo "64GB eMMC detected"
    sed -i 's/compress=zstd:1/noatime,lazytime,discard=sync,compress-force=zstd:3,space_cache=v2/g' /etc/fstab
  else
    sed -i 's/compress=zstd:1/noatime,lazytime,commit=120,discard=async,compress-force=zstd:1,space_cache=v2/g' /etc/fstab
  fi
else
  echo "No fstab param adjustments needed"
fi

# DECK KERNEL ARGUMENTS
if [[ ":Jupiter:" =~ ":$SYS_ID:"  ]]; then
  echo "Checking for needed karg changes (Jupiter)"

  if [[ ! $KARGS =~ "amd_pstate" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=amd_pstate=active"
  fi

  if [[ ! $KARGS =~ "amd_iommu" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=amd_iommu=off"
  fi

  if [[ ! $KARGS =~ "amdgpu.gttsize" ]]; then
    if [[ "$(awk '/MemTotal/{print $(NF-1)}' /proc/meminfo)" == "31664740" ]]; then
      echo "32GB RAM Steam Deck detected"
      NEEDED_KARGS="$NEEDED_KARGS --append=amdgpu.gttsize=16256"
    else
      NEEDED_KARGS="$NEEDED_KARGS --append=amdgpu.gttsize=8128"
    fi
  fi

  if [[ ! $KARGS =~ "spi_amd.speed_dev" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=spi_amd.speed_dev=1"
  fi

  if [[ ! $KARGS =~ "initcall_blacklist" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=initcall_blacklist=simpledrm_platform_driver_init"
  fi
elif [[ $IMAGE_NAME =~ "deck"  ]]; then
  echo "Generic device detected. Performing setup..."

  systemctl disable --now jupiter-fan-control.service
  systemctl disable --now vpower.service
  systemctl disable --now jupiter-biosupdate.service
  systemctl disable --now jupiter-controller-update.service
  systemctl disable --now ryzenadj.service
  systemctl disable --now batterylimit.service
  systemctl --global disable --now sdgyrodsu.service
fi

if grep -qz "Kernel driver in use: radeon" <<< $GPU_ID; then
  echo "Legacy AMD hardware detected, enabling CIK and SI support in AMDGPU"
  if [[ ! $KARGS =~ "radeon.si_support" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=radeon.si_support=0"
  fi

  if [[ ! $KARGS =~ "radeon.cik_support" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=radeon.cik_support=0"
  fi

  if [[ ! $KARGS =~ "amdgpu.si_support" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=amdgpu.si_support=1"
  fi

  if [[ ! $KARGS =~ "amdgpu.cik_support" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=amdgpu.cik_support=1"
  fi
fi

if [[ $IMAGE_FLAVOR = "nvidia"  ]]; then
  echo "Checking for needed karg changes (Nvidia)"

  if [[ ! $KARGS =~ "rd.driver.blacklist" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=rd.driver.blacklist=nouveau"
  fi

  if [[ ! $KARGS =~ "modprobe.blacklist" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=modprobe.blacklist=nouveau"
  fi

  if [[ ! $KARGS =~ "nvidia-drm.modeset" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=nvidia-drm.modeset=1"
  fi
else
  echo "Checking for needed karg changes"

  if [[ $KARGS =~ "rd.driver.blacklist" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --delete=rd.driver.blacklist=nouveau"
  fi

  if [[ $KARGS =~ "modprobe.blacklist" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --delete=modprobe.blacklist=nouveau"
  fi

  if [[ $KARGS =~ "nvidia-drm.modeset" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --delete=nvidia-drm.modeset=1"
  fi
fi

if [[ $KARGS =~ "nomodeset" ]]; then
  echo "Removing nomodeset"
  NEEDED_KARGS="$NEEDED_KARGS --delete=nomodeset"
fi

if [[ ! $KARGS =~ "rd.luks.options" ]]; then
  NEEDED_KARGS="$NEEDED_KARGS --append=rd.luks.options=discard"
fi

if [[ -n "$NEEDED_KARGS" ]]; then
  echo "Found needed karg changes, applying the following: $NEEDED_KARGS"
  rpm-ostree kargs ${NEEDED_KARGS} --reboot || exit 1
else
  echo "No karg changes needed"
fi

echo $HWS_VER > $HWS_VER_FILE
echo $IMAGE_NAME > $KNOWN_IMAGE_NAME_FILE
echo $IMAGE_FLAVOR > $KNOWN_IMAGE_FLAVOR_FILE