#!/usr/bin/env bash
# Configure kernel boot parameters for unified memory optimization
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "$SCRIPT_DIR/../../lib/common.sh"
source "$SCRIPT_DIR/../../lib/detect.sh"
source "$SCRIPT_DIR/../../lib/format.sh"

GRUB_FILE="/etc/default/grub"

log_header "Kernel Boot Parameter Optimization"

# ── Check root early ────────────────────────────────────
if [[ $EUID -ne 0 ]]; then
    log_error "This script requires root. Re-run with: sudo make optimize-kernel"
    exit 1
fi

# ── Show current state ───────────────────────────────────
log_info "Current kernel command line:"
printf "  ${DIM}%s${RESET}\n" "$(cat /proc/cmdline)"
echo ""

param_iommu="$(detect_kernel_param 'iommu')"
param_gttsize="$(detect_gttsize_param)"
param_pages="$(detect_pages_limit_param)"

rec_gttsize="$(recommended_gttsize_mib)"
rec_pages="$(recommended_pages_limit)"

# ── Check what's needed ──────────────────────────────────
needs_change=false

log_info "Parameter status:"

if [[ "$param_iommu" == "pt" ]]; then
    print_status pass "iommu=pt" "already set"
else
    print_status fail "iommu=pt" "$([ -n "$param_iommu" ] && echo "current: $param_iommu" || echo "missing")"
    needs_change=true
fi

if [[ -n "$param_gttsize" ]] && (( param_gttsize >= rec_gttsize )); then
    print_status pass "amdgpu.gttsize" "current: $param_gttsize MiB"
else
    print_status fail "amdgpu.gttsize" "$([ -n "$param_gttsize" ] && echo "current: $param_gttsize MiB, " || echo "missing, ")recommended: $rec_gttsize MiB (~$(human_mib "$rec_gttsize"))"
    needs_change=true
fi

if [[ -n "$param_pages" ]] && (( param_pages >= rec_pages )); then
    print_status pass "ttm.pages_limit" "current: $param_pages"
else
    print_status fail "ttm.pages_limit" "$([ -n "$param_pages" ] && echo "current: $param_pages, " || echo "missing, ")recommended: $rec_pages"
    needs_change=true
fi

if ! $needs_change; then
    echo ""
    log_success "All kernel parameters are already optimal!"
    exit 0
fi

# ── Explain what we're doing ─────────────────────────────
echo ""
log_info "These parameters enable unified memory for the integrated GPU:"
echo "  iommu=pt              IOMMU passthrough — reduces memory access latency"
echo "  amdgpu.gttsize=$rec_gttsize  GPU can dynamically access ~$(human_mib "$rec_gttsize") system RAM"
echo "  ttm.pages_limit=$rec_pages  Pin limit for GPU memory pages ($(human_mib "$rec_gttsize") in 4K pages)"
echo ""

# ── Apply changes ────────────────────────────────────────
if ! confirm "Apply these kernel parameters to GRUB?"; then
    log_info "Skipped. You can apply manually by editing $GRUB_FILE"
    exit 0
fi

# Backup
BACKUP_DIR="$(data_dir backups)"
backup_file="$BACKUP_DIR/grub-$(timestamp).bak"
cp "$GRUB_FILE" "$backup_file"
log_success "GRUB backup saved: $backup_file"

# Parse current GRUB_CMDLINE_LINUX using Python (data via env vars, not interpolation)
current_cmdline="$(GRUB_PATH="$GRUB_FILE" python3 -c '
import re, os
with open(os.environ["GRUB_PATH"]) as f:
    for line in f:
        m = re.match(r"^GRUB_CMDLINE_LINUX=\"(.*)\"", line)
        if m:
            print(m.group(1))
            raise SystemExit(0)
print("")
')"

# Remove any existing values of these params
new_cmdline="$current_cmdline"
new_cmdline="$(echo "$new_cmdline" | sed -E 's/\biommu=[^ ]*//g')"
new_cmdline="$(echo "$new_cmdline" | sed -E 's/\bamd_iommu=[^ ]*//g')"
new_cmdline="$(echo "$new_cmdline" | sed -E 's/\bamdgpu\.gttsize=[^ ]*//g')"
new_cmdline="$(echo "$new_cmdline" | sed -E 's/\bttm\.pages_limit=[^ ]*//g')"
# Clean up extra spaces
new_cmdline="$(echo "$new_cmdline" | xargs)"

# Add new params
new_cmdline="$new_cmdline iommu=pt amdgpu.gttsize=$rec_gttsize ttm.pages_limit=$rec_pages"

log_info "GRUB_CMDLINE_LINUX change:"
printf "  ${RED}Before:${RESET} %s\n" "$current_cmdline"
printf "  ${GREEN}After:${RESET}  %s\n" "$new_cmdline"
echo ""

if ! confirm "Write this change?"; then
    log_info "Aborted. Backup remains at: $backup_file"
    exit 0
fi

# Apply using Python (all data via env vars — no shell interpolation into Python code)
GRUB_PATH="$GRUB_FILE" NEW_CMDLINE="$new_cmdline" python3 -c '
import re, os
grub_path = os.environ["GRUB_PATH"]
new_line = "GRUB_CMDLINE_LINUX=\"" + os.environ["NEW_CMDLINE"] + "\""
with open(grub_path) as f:
    content = f.read()
content = re.sub(r"^GRUB_CMDLINE_LINUX=.*", new_line, content, count=1, flags=re.MULTILINE)
with open(grub_path, "w") as f:
    f.write(content)
'
log_success "GRUB config updated"

# Regenerate GRUB — prefer grubby on modern Fedora (BLS), fall back to grub2-mkconfig
log_info "Regenerating boot configuration..."
if is_cmd grubby; then
    grubby --update-kernel=ALL --args="iommu=pt amdgpu.gttsize=$rec_gttsize ttm.pages_limit=$rec_pages"
    log_success "Boot entries updated via grubby"
elif [[ -d /boot/grub2 ]]; then
    grub2-mkconfig -o /boot/grub2/grub.cfg
    log_success "GRUB regenerated via grub2-mkconfig"
elif [[ -d /boot/grub ]]; then
    grub-mkconfig -o /boot/grub/grub.cfg
    log_success "GRUB regenerated via grub-mkconfig"
else
    log_error "Could not find grubby or grub config directory. Regenerate manually."
    exit 1
fi

echo ""
log_warn "REBOOT REQUIRED for kernel parameters to take effect."
log_info "After reboot, verify with: make audit"