#!/bin/bash
# =============================================================================
# Step 2: Configure System Environment — AMD BC-250 ROCm
# =============================================================================
# Sets up:
#   1. User group membership (render, video)
#   2. Environment variables (fish shell + bash fallback)
#   3. Modprobe configuration (amdgpu.conf)
#   4. Boot parameters (Limine or GRUB)
#
# Run as: bash 02_configure_environment.sh
# Some steps will prompt for sudo.
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
GUIDE_DIR="$(dirname "$SCRIPT_DIR")"
CURRENT_USER="$(whoami)"

echo "============================================"
echo " Configuring System for BC-250 ROCm"
echo "============================================"

# ─────────────────────────────────────────────
# 1. User Groups
# ─────────────────────────────────────────────
echo ""
echo "[1/4] Adding $CURRENT_USER to render and video groups..."
sudo usermod -aG render,video "$CURRENT_USER"
echo "  Done. (Requires re-login to take effect)"

# ─────────────────────────────────────────────
# 2. Environment Variables
# ─────────────────────────────────────────────
echo ""
echo "[2/4] Configuring environment variables..."

# Detect shell
if command -v fish &>/dev/null && [[ "$SHELL" == *fish* ]]; then
    echo "  Detected fish shell — configuring fish universal vars"
    
    fish -c '
        set -Ux HSA_OVERRIDE_GFX_VERSION 10.1.0
        set -Ux HIP_VISIBLE_DEVICES 0
        set -Ux HSA_ENABLE_SDMA 0
        set -Ux HSA_TOOLS_LIB ""
        set -Ux HSA_TOOLS_REPORT_LOAD_FAILURE 0
        set -Ux ROCM_PATH /opt/rocm
        set -Ux HIP_PATH /opt/rocm
        echo "  Fish universal variables set."
    '
    
    # Also add to fish config for PATH
    FISH_CONFIG="$HOME/.config/fish/config.fish"
    if ! grep -q "ROCm PATH" "$FISH_CONFIG" 2>/dev/null; then
        mkdir -p "$(dirname "$FISH_CONFIG")"
        cat >> "$FISH_CONFIG" << 'FISHEOF'

# ROCm PATH
if not contains /opt/rocm/bin $PATH
    set -gx PATH /opt/rocm/bin $PATH
end
FISHEOF
        echo "  Added /opt/rocm/bin to fish PATH"
    fi
else
    echo "  Configuring bash/zsh environment..."
    
    PROFILE="$HOME/.bashrc"
    if ! grep -q "HSA_OVERRIDE_GFX_VERSION" "$PROFILE" 2>/dev/null; then
        cat >> "$PROFILE" << 'BASHEOF'

# ===== ROCm BC-250 Environment =====
export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HIP_VISIBLE_DEVICES=0
export HSA_ENABLE_SDMA=0
export HSA_TOOLS_LIB=""
export HSA_TOOLS_REPORT_LOAD_FAILURE=0
export ROCM_PATH=/opt/rocm
export HIP_PATH=/opt/rocm
export PATH=/opt/rocm/bin:$PATH
# ===== End ROCm =====
BASHEOF
        echo "  Added to $PROFILE"
    else
        echo "  Already configured in $PROFILE"
    fi
fi

echo ""
echo "  CRITICAL BC-250 Environment Variables:"
echo "    HSA_OVERRIDE_GFX_VERSION=10.1.0  (map gfx1013 → gfx1010)"
echo "    HIP_VISIBLE_DEVICES=0            (single GPU)"
echo "    HSA_ENABLE_SDMA=0                (SDMA broken on Cyan Skillfish)"
echo "    HSA_TOOLS_LIB=\"\"                 (disable profiling tools)"
echo ""
echo "  DO NOT SET these (they cause hangs):"
echo "    HIP_LAUNCH_BLOCKING=1"
echo "    GPU_MAX_HW_QUEUES=1"

# ─────────────────────────────────────────────
# 3. Modprobe Configuration
# ─────────────────────────────────────────────
echo ""
echo "[3/4] Installing modprobe configuration..."

AMDGPU_CONF="$GUIDE_DIR/config/amdgpu.conf"
if [[ -f "$AMDGPU_CONF" ]]; then
    sudo cp "$AMDGPU_CONF" /etc/modprobe.d/amdgpu.conf
    echo "  Installed /etc/modprobe.d/amdgpu.conf"
    echo "  Contents: $(cat /etc/modprobe.d/amdgpu.conf)"
else
    echo "  Creating amdgpu.conf manually..."
    echo 'options amdgpu noretry=0 gpu_recovery=1 sched_hw_submission=2 ppfeaturemask=0xfff73ef7' | \
        sudo tee /etc/modprobe.d/amdgpu.conf
fi

# ─────────────────────────────────────────────
# 4. Boot Parameters
# ─────────────────────────────────────────────
echo ""
echo "[4/4] Configuring boot parameters..."

BOOT_PARAMS="amdgpu.ppfeaturemask=0xfff73ef7 amdgpu.noretry=0 amdgpu.gpu_recovery=1 amdgpu.sched_hw_submission=2"

if [[ -f /boot/limine.conf ]]; then
    echo "  Detected Limine bootloader"
    if ! grep -q "ppfeaturemask" /boot/limine.conf; then
        echo ""
        echo "  Add these parameters to your kernel cmdline in /boot/limine.conf:"
        echo "    $BOOT_PARAMS"
        echo ""
        echo "  Example: Find the 'cmdline=' line and append the parameters."
        echo "  Then run: sudo limine-update"
    else
        echo "  Boot parameters already configured in /boot/limine.conf"
    fi
elif [[ -f /etc/default/grub ]]; then
    echo "  Detected GRUB bootloader"
    if ! grep -q "ppfeaturemask" /etc/default/grub; then
        sudo sed -i "s|GRUB_CMDLINE_LINUX_DEFAULT=\"|GRUB_CMDLINE_LINUX_DEFAULT=\"$BOOT_PARAMS |" /etc/default/grub
        sudo grub-mkconfig -o /boot/grub/grub.cfg
        echo "  GRUB updated with boot parameters"
    else
        echo "  Boot parameters already configured in GRUB"
    fi
else
    echo "  WARNING: Unknown bootloader. Manually add these kernel parameters:"
    echo "    $BOOT_PARAMS"
fi

echo ""
echo "============================================"
echo " System configuration complete"
echo "============================================"
echo ""
echo "NOTE: Log out and back in for group changes to take effect."
echo "      Boot parameters require a reboot after module installation."
echo ""
echo "Next: Download kernel source and build patched module"
echo "  Run: bash 04_build_patched_module.sh"
