#!/usr/bin/env bash

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

# SCRIPT VERSION
HWS_VER=9
HWS_VER_FILE="/etc/bazzite/hws_version"
HWS_VER_RAN=$(cat $HWS_VER_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)
KNOWN_FEDORA_VERSION_FILE="/etc/bazzite/fedora_version"
KNOWN_FEDORA_VERSION=$(cat $KNOWN_FEDORA_VERSION_FILE)

# 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)
INITRAMFS=$(rpm-ostree initramfs)
NEEDED_KARGS=""
echo "Current kargs: $KARGS"
mkdir -p /etc/bazzite

# 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"
      sed -i 's/zram-size=1024/zram-size=2048/g' /etc/systemd/zram-generator.conf

    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

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

  if [[ ! $KARGS =~ "nmi_watchdog" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=nmi_watchdog=0"
  fi
  if [[ ! $KARGS =~ "modprobe.blacklist=sp5100_tco" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=modprobe.blacklist=sp5100_tco"
  fi
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=nouveau" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --append=rd.driver.blacklist=nouveau"
  fi

  if [[ ! $KARGS =~ "modprobe.blacklist=nouveau" ]]; 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=nouveau" ]]; then
    NEEDED_KARGS="$NEEDED_KARGS --delete-if-present=rd.driver.blacklist=nouveau"
  fi

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

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

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

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

if [[ $INITRAMFS =~ "disabled" ]]; then
  echo "Found needed initramfs changes, applying."
  plymouth display-message --text="Updating initramfs - Please wait" || true
  rpm-ostree initramfs --enable --arg=-I --arg=/etc/crypttab
fi

if [[ -n "$NEEDED_KARGS" ]]; then
  echo "Found needed karg changes, applying the following: $NEEDED_KARGS"
  plymouth display-message --text="Updating kargs - Please wait" || true
  rpm-ostree kargs ${NEEDED_KARGS} --reboot || exit 1
else
  echo "No karg changes needed"
fi

# RUN REMAINDER OF SCRIPT ONLY 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
    if [[ -f $KNOWN_FEDORA_VERSION_FILE && -f $KNOWN_FEDORA_VERSION_FILE ]]; then
      # Run script if image has been rebased
      if [[ $IMAGE_NAME = $KNOWN_IMAGE_NAME && $IMAGE_FLAVOR = $KNOWN_IMAGE_FLAVOR && $FEDORA_VERSION = $KNOWN_FEDORA_VERSION ]]; then
        echo "Hardware setup has already run. Exiting..."
        exit 0
      fi
    fi
  fi
fi

if [[ $IMAGE_NAME =~ "deck" || $IMAGE_NAME =~ "ally" || $IMAGE_NAME =~ "framegame" ]]; then
  if [[ ":Jupiter:" =~ ":$SYS_ID:" ]]; then
    # Future updates to Deck on Jupiter here
    echo "Jupiter hardware detected, skipping further setup..."
  else
    echo "Generic device detected. Performing setup..."
    if [[ ":ROG Ally RC71L_RC71L:AYANEO GEEK:AYANEO 2:AYANEO 2S:AOKZOE A1 AR07:G1619-04:" =~ ":$SYS_ID:" ]]; then
      echo "HandyGCCS supported handheld detected, enabling handycon & TDP control..."
      systemctl enable --now handycon.service
      sed -i 's/ENABLE_HARDWARE_CONTROL_ON_NON_DECK_HARDWARE=0/ENABLE_HARDWARE_CONTROL_ON_NON_DECK_HARDWARE=1/g' /etc/default/steam-hardware-control
    fi
    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
fi

# 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

# HOSTNAME FIX
# If the hostname is too long Distrobox will fail during setup
# Additonally, Anaconda likes to set the hostname to be the ipv6 address
# Let's check the length and reset it to something sensible if that happens.
if hostname | wc -m > 15; then
  hostnamectl set-hostname bazzite
fi

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