#!/bin/bash
# =============================================================================
# Step 4: Download CachyOS Kernel Source, Patch & Build amdgpu Module
# =============================================================================
# CRITICAL: Must use CachyOS kernel source (NOT vanilla kernel.org source).
#           Vanilla source has different struct layouts → NULL pointer crash.
#
# This script:
#   1. Detects current kernel version
#   2. Downloads matching CachyOS source PKGBUILD
#   3. Extracts and prepares the source tree
#   4. Applies 3 BC-250 patches (GFXOFF, KIQ bypass, dead-GPU detection)
#   5. Builds ONLY the amdgpu.ko module (not the entire kernel)
#   6. Strips debug symbols and compresses with zstd
#
# Run as: bash 04_build_patched_module.sh [build_dir]
#   Default build_dir: ~/kernel-build
#
# Prerequisites:
#   - CachyOS with ROCm packages installed
#   - base-devel, asp/pkgctl, clang, llvm, lld installed
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
GUIDE_DIR="$(dirname "$SCRIPT_DIR")"
PATCHES_DIR="$GUIDE_DIR/patches"
BUILD_DIR="${1:-$HOME/kernel-build}"
KERNEL_VER="$(uname -r)"

echo "============================================"
echo " Building Patched amdgpu Module for BC-250"
echo "============================================"
echo ""
echo "  Kernel:    $KERNEL_VER"
echo "  Build dir: $BUILD_DIR"
echo "  Patches:   $PATCHES_DIR"
echo ""

# ─────────────────────────────────────────────
# 1. Install Build Dependencies
# ─────────────────────────────────────────────
echo "[1/7] Installing build dependencies..."
sudo pacman -S --needed --noconfirm base-devel asp clang llvm lld bc cpio

# ─────────────────────────────────────────────
# 2. Create Build Directory
# ─────────────────────────────────────────────
echo "[2/7] Creating build directory..."
mkdir -p "$BUILD_DIR"
cd "$BUILD_DIR"

# ─────────────────────────────────────────────
# 3. Get CachyOS Kernel Source
# ─────────────────────────────────────────────
echo "[3/7] Getting CachyOS kernel source..."

# Extract base version from running kernel (e.g., 6.19.6-2-cachyos → 6.19.6)
BASE_VER=$(echo "$KERNEL_VER" | sed 's/-.*//')
echo "  Base kernel version: $BASE_VER"

# Try to get the PKGBUILD via asp (Arch Source Package)
if command -v pkgctl &>/dev/null; then
    echo "  Using pkgctl to fetch source..."
    pkgctl repo clone --protocol=https linux-cachyos 2>/dev/null || true
elif command -v asp &>/dev/null; then
    echo "  Using asp to fetch source..."
    asp checkout linux-cachyos 2>/dev/null || true
fi

# If asp/pkgctl didn't work, try direct download
if [[ ! -d linux-cachyos ]]; then
    echo "  Trying CachyOS GitHub..."
    if command -v git &>/dev/null; then
        git clone --depth 1 https://github.com/CachyOS/linux-cachyos.git 2>/dev/null || true
    fi
fi

if [[ ! -d linux-cachyos ]]; then
    echo ""
    echo "  ERROR: Could not fetch CachyOS kernel source automatically."
    echo "  Please manually download the CachyOS kernel source matching $KERNEL_VER"
    echo "  and place it in $BUILD_DIR/linux-cachyos/"
    echo ""
    echo "  Options:"
    echo "    1. pkgctl repo clone --protocol=https linux-cachyos"
    echo "    2. git clone https://github.com/CachyOS/linux-cachyos.git"
    echo "    3. Download from https://github.com/CachyOS/linux-cachyos/releases"
    exit 1
fi

cd linux-cachyos

# Run makepkg to download and extract source (don't build yet)
echo "  Extracting source (makepkg --nobuild)..."
makepkg --nobuild --skippgpcheck 2>&1 | tail -5

# Find the extracted source directory
SRC_DIR=$(find . -maxdepth 2 -name "drivers" -type d | head -1 | sed 's|/drivers$||')
if [[ -z "$SRC_DIR" ]]; then
    SRC_DIR=$(find "$BUILD_DIR" -maxdepth 3 -name "Kconfig" -path "*/linux-*" | head -1 | xargs dirname 2>/dev/null || true)
fi

if [[ -z "$SRC_DIR" || ! -d "$SRC_DIR" ]]; then
    echo "  ERROR: Could not find extracted kernel source."
    echo "  Look for a linux-* directory under $BUILD_DIR/linux-cachyos/src/"
    exit 1
fi

echo "  Source dir: $SRC_DIR"
AMDGPU_DIR="$SRC_DIR/drivers/gpu/drm/amd/amdgpu"

if [[ ! -d "$AMDGPU_DIR" ]]; then
    echo "  ERROR: amdgpu source not found at $AMDGPU_DIR"
    exit 1
fi

# ─────────────────────────────────────────────
# 4. Prepare Build Environment
# ─────────────────────────────────────────────
echo "[4/7] Preparing build environment..."
cd "$SRC_DIR"

# Copy running kernel's config
zcat /proc/config.gz > .config 2>/dev/null || cp /boot/config-"$KERNEL_VER" .config 2>/dev/null || true

if [[ ! -f .config ]]; then
    echo "  ERROR: Could not find kernel config"
    exit 1
fi

# Prepare the build environment
make LLVM=1 olddefconfig
make LLVM=1 modules_prepare

# ─────────────────────────────────────────────
# 5. Apply BC-250 Patches
# ─────────────────────────────────────────────
echo "[5/7] Applying BC-250 patches..."

echo "  Applying patch 1: GFXOFF disable (gfx_v10_0.c)..."
python3 "$PATCHES_DIR/patch1_gfxoff.py" "$AMDGPU_DIR"

echo "  Applying patch 2: KIQ bypass + dead-GPU (gmc_v10_0.c)..."
python3 "$PATCHES_DIR/patch2_gmc.py" "$AMDGPU_DIR"

echo "  Applying patch 3: KIQ bypass + dead-GPU (amdgpu_gmc.c)..."
python3 "$PATCHES_DIR/patch3_amdgpu_gmc.py" "$AMDGPU_DIR"

echo "  All patches applied."

# ─────────────────────────────────────────────
# 6. Build the amdgpu Module
# ─────────────────────────────────────────────
echo "[6/7] Building amdgpu module (LLVM=1)..."
echo "  This may take several minutes..."

# Get number of CPU cores for parallel build
NPROC=$(nproc)
echo "  Using $NPROC parallel jobs"

make LLVM=1 -j"$NPROC" M=drivers/gpu/drm/amd/amdgpu modules 2>&1 | tail -20

# Check build result
MODULE="drivers/gpu/drm/amd/amdgpu/amdgpu.ko"
if [[ ! -f "$MODULE" ]]; then
    echo "  ERROR: Build failed — amdgpu.ko not found"
    echo "  Check the build output above for errors."
    exit 1
fi

echo "  Build successful!"
ls -lh "$MODULE"

# ─────────────────────────────────────────────
# 7. Strip and Compress
# ─────────────────────────────────────────────
echo "[7/7] Stripping debug symbols and compressing..."

strip --strip-debug "$MODULE"
zstd -f --rm "$MODULE" -o "${MODULE}.zst"

FINAL_MODULE="${MODULE}.zst"
echo "  Final module: $FINAL_MODULE"
ls -lh "$FINAL_MODULE"

echo ""
echo "============================================"
echo " Module built successfully!"
echo "============================================"
echo ""
echo " Module: $(realpath "$FINAL_MODULE")"
echo " Size:   $(du -h "$FINAL_MODULE" | cut -f1)"
echo ""
echo "Next: Run 05_install_module.sh to install and reboot"
echo "  bash $SCRIPT_DIR/05_install_module.sh $(realpath "$FINAL_MODULE")"