#!/usr/bin/bash
set -euo pipefail

# ============================================================
# post-upgrade.sh — Bazzite post-upgrade automation
# Run this after every major Bazzite upgrade to restore
# custom packages, services, plugins, and tweaks.
#
# Usage:
#   chmod +x ~/post-upgrade.sh
#   sudo ~/post-upgrade.sh
# ============================================================

RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m'

log()  { echo -e "${GREEN}[+]${NC} $*"; }
warn() { echo -e "${YELLOW}[!]${NC} $*"; }
err()  { echo -e "${RED}[x]${NC} $*"; }

NEED_REBOOT=0

reboot_needed() {
  warn "$* — reboot will be needed"
  NEED_REBOOT=1
}

# ── sysfs helpers ────────────────────────────────────────
write_sysfs() {
  local path="$1" val="$2"
  if [[ -w "$path" ]]; then
    echo "$val" > "$path" 2>/dev/null || warn "Failed to write $val to $path"
  fi
}

# ── 1. rpm-ostree layered packages ───────────────────────
log "=== Step 1: Layered packages ==="

LAYER_PKGS=(
  coolercontrol
  liquidctl
  solaar
  mangohud
  lm_sensors
)

for pkg in "${LAYER_PKGS[@]}"; do
  if rpm-ostree status --json 2>/dev/null | grep -q "\"$pkg\""; then
    log "  $pkg already layered"
  else
    log "  Layering $pkg ..."
    rpm-ostree install -y "$pkg" && reboot_needed "rpm-ostree layered $pkg"
  fi
done

# ── 2. Kernel args ───────────────────────────────────────
log "=== Step 2: Kernel arguments ==="

if rpm-ostree kargs | grep -q "mitigations=off"; then
  log "  mitigations=off already set"
else
  log "  Adding mitigations=off ..."
  rpm-ostree kargs --append-if-missing="mitigations=off" && reboot_needed "kernel arg mitigations=off"
fi

# ── 2b. Kernel args: disable ertm (bluetooth) ────────────
if rpm-ostree kargs | grep -q "bluetooth.disable_ertm=1"; then
  log "  bluetooth.disable_ertm=1 already set"
else
  log "  Adding bluetooth.disable_ertm=1 ..."
  rpm-ostree kargs --append-if-missing="bluetooth.disable_ertm=1" && reboot_needed "kernel arg bluetooth.disable_ertm=1"
fi

# ── 3. Modules ───────────────────────────────────────────
log "=== Step 3: Kernel modules ==="

MODULES_DIR="/etc/modules-load.d"
mkdir -p "$MODULES_DIR"

# nct6687 — sensor chip
if [[ -f "$MODULES_DIR/nct6687.conf" ]]; then
  log "  nct6687 module already configured"
else
  log "  Configuring nct6687 module ..."
  echo 'nct6687' > "$MODULES_DIR/nct6687.conf"
fi

# ── 4. Oberon GPU governor ───────────────────────────────
log "=== Step 4: Oberon GPU frequency governor ==="

if systemctl is-enabled oberon-governor.service &>/dev/null; then
  log "  oberon-governor already installed and enabled"
else
  log "  Installing oberon-governor ..."
  curl -s https://raw.githubusercontent.com/vietsman/bc250-documentation/refs/heads/main/oberon-setup.sh | sh
  systemctl enable oberon-governor
  systemctl start oberon-governor || warn "oberon-governor failed to start (may need reboot)"
fi

# ── 5. Udev rules ────────────────────────────────────────
log "=== Step 5: udev rules ==="

UDEV_RULE="/etc/udev/rules.d/99-ryzenadj-power-source-change.rules"
if [[ -f "$UDEV_RULE" ]]; then
  log "  ryzenadj udev rule already present"
else
  warn "  ryzenadj udev rule missing — this was previously installed"
  warn "  Check source / reinstall via ujust enable-ryzenadj-max-performance"
fi

# ── 6. Services ──────────────────────────────────────────
log "=== Step 6: Services ==="

for svc in coolercontrold sshd; do
  if systemctl is-enabled "$svc.service" &>/dev/null; then
    log "  $svc already enabled"
  else
    systemctl enable --now "$svc.service" && log "  $svc enabled and started"
  fi
done

# ── 7. Firewall ──────────────────────────────────────────
log "=== Step 7: Firewall ==="

if firewall-cmd --list-services 2>/dev/null | grep -q ssh; then
  log "  SSH already allowed through firewall"
else
  log "  Opening SSH port ..."
  firewall-cmd --permanent --add-service=ssh
  firewall-cmd --reload
fi

# ── 8. ujust tweaks ──────────────────────────────────────
log "=== Step 8: ujust performance tweaks ==="

# Decky plugin loader
if [[ -d "$HOME/homebrew" ]]; then
  log "  Decky Loader already present (in /home, survives upgrades)"
else
  log "  Installing Decky Loader ..."
  ujust setup-decky install
fi

log "  Running ujust enable-ryzenadj-max-performance ..."
ujust enable-ryzenadj-max-performance 2>/dev/null || warn "ujust enable-ryzenadj-max-performance failed (may already be applied)"

log "  Running ujust get-decky-bazzite-buddy ..."
ujust get-decky-bazzite-buddy 2>/dev/null || warn "ujust get-decky-bazzite-buddy failed (may already be installed)"

# ── 9. SimpleDeckyTDP plugin ──────────────────────────────
log "=== Step 9: SimpleDeckyTDP Decky plugin ==="

if [[ -d "$HOME/homebrew/plugins/SimpleDeckyTDP" ]]; then
  log "  SimpleDeckyTDP already installed"
else
  log "  Installing SimpleDeckyTDP ..."
  rm -rf "$HOME/homebrew/plugins/SimpleDeckyTDP"
  TDP_URL=$(curl -s https://api.github.com/repos/aarron-lee/SimpleDeckyTDP/releases/latest \
    | grep "browser_download_url" | cut -d '"' -f 4)
  if [[ -n "$TDP_URL" ]]; then
    curl -L "$TDP_URL" -o /tmp/SimpleDeckyTDP.zip
    7z x /tmp/SimpleDeckyTDP.zip -o"$HOME/homebrew/plugins"
    rm /tmp/SimpleDeckyTDP.zip
    systemctl restart plugin_loader.service
  else
    warn "  Could not determine SimpleDeckyTDP download URL"
  fi
fi

# ── 10. SSH authorized keys ──────────────────────────────
log "=== Step 10: SSH keys ==="

mkdir -p "$HOME/.ssh"
chmod 700 "$HOME/.ssh"

if grep -q "cachyos-otter" "$HOME/.ssh/authorized_keys" 2>/dev/null; then
  log "  SSH key already present"
else
  echo "ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIHtbk6IRSDWGCiET+7NQ/BTWqF4VQtb3T8EecV2pEWhU cachyos-otter" >> "$HOME/.ssh/authorized_keys"
  chmod 600 "$HOME/.ssh/authorized_keys"
fi

# ── Finish ────────────────────────────────────────────────
echo ""
log "=== Done ==="

if [[ $NEED_REBOOT -eq 1 ]]; then
  echo ""
  warn "REBOOT REQUIRED for rpm-ostree layers / kernel args to take effect"
  echo "  Run: systemctl reboot"
else
  log "No reboot needed. All customizations are applied."
fi

log "If oberon-governor isn't running, check /etc/oberon-config.yaml"