# shellcheck shell=bash
# Z13 unified memory profile helpers (GZ302 / Strix Halo)

set -euo pipefail

INSTALL_ROOT="${INSTALL_ROOT:-$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)}"
CONFIG_FILE="${CONFIG_FILE:-${INSTALL_ROOT}/config/profiles.conf}"
TTM_CONF="/etc/modprobe.d/ttm.conf"
TTM_PARAM="/sys/module/ttm/parameters/pages_limit"
MAX_TTM_PERCENT=90

log() { printf '[z13-mem] %s\n' "$*" >&2; }

require_root() {
    if [[ "${EUID}" -ne 0 ]]; then
        echo "Root required. Re-run with pkexec or sudo." >&2
        exit 1
    fi
}

total_ram_mb() {
    awk '/^MemTotal:/ { printf "%d\n", $2 / 1024 }' /proc/meminfo
}

total_ram_gb() {
    awk '/^MemTotal:/ { printf "%.1f\n", $2 / 1024 / 1024 }' /proc/meminfo
}

gb_to_pages() {
    python3 -c "print(int(float('${1}') * (1024 ** 3) / 4096))"
}

pages_to_gb() {
    python3 -c "print(f'{int(${1}) * 4096 / (1024 ** 3):.2f}')"
}

find_uma_base() {
    local card
    for card in /sys/class/drm/card*/device; do
        [[ -d "${card}/uma" ]] || continue
        if [[ -r "${card}/uma/carveout_options" ]]; then
            echo "${card}/uma"
            return 0
        fi
    done
    return 1
}

parse_carveout_options() {
    local uma_base="$1"
    cat "${uma_base}/carveout_options"
}

# Pick carveout index closest to target_mb without exceeding max_mb (when set).
find_carveout_index() {
    local uma_base="$1" target_mb="$2" max_mb="${3:-}"
    python3 - "${uma_base}/carveout_options" "${target_mb}" "${max_mb}" <<'PY'
import re, sys
path, target_s, max_s = sys.argv[1:4]
target = float(target_s)
max_mb = float(max_s) if max_s else None
options = []
with open(path, encoding="utf-8", errors="replace") as f:
    for line in f:
        m = re.match(r"\s*(\d+):\s*(.*)$", line.strip())
        if not m:
            continue
        idx, rest = int(m.group(1)), m.group(2).lower()
        mb = None
        g = re.search(r"(\d+(?:\.\d+)?)\s*gb", rest)
        if g:
            mb = float(g.group(1)) * 1024
        else:
            m2 = re.search(r"(\d+(?:\.\d+)?)\s*mb", rest)
            if m2:
                mb = float(m2.group(1))
        if mb is None:
            continue
        if max_mb is not None and mb > max_mb:
            continue
        options.append((idx, mb, line.strip()))
if not options:
    sys.exit(1)
# Prefer smallest option >= target; else largest below target.
at_or_above = [o for o in options if o[1] >= target - 0.5]
if at_or_above:
    best = min(at_or_above, key=lambda o: o[1])
else:
    best = max(options, key=lambda o: o[1])
print(best[0])
print(best[2])
PY
}

read_profile_value() {
    local section="$1" key="$2"
    awk -v section="[$section]" -v key="$key" '
        $0 == section { found=1; next }
        /^\[/ { found=0 }
        found && $0 ~ "^[[:space:]]*" key "=" {
            sub(/^[^=]*=/, "")
            gsub(/^[[:space:]]+|[[:space:]]+$/, "")
            print
            exit
        }
    ' "${CONFIG_FILE}"
}

current_ttm_gb() {
    if [[ ! -r "${TTM_PARAM}" ]]; then
        echo "unknown"
        return
    fi
    pages_to_gb "$(cat "${TTM_PARAM}")"
}

current_carveout_info() {
    local uma_base
    if ! uma_base="$(find_uma_base)"; then
        echo "UMA sysfs: not available"
        return
    fi
    echo "Current index: $(cat "${uma_base}/carveout")"
    echo "Options:"
    parse_carveout_options "${uma_base}"
}

validate_profile() {
    local profile="$1"
    local carveout_mb ttm_action ttm_gb min_reserve total_mb total_gb max_ttm_gb

    carveout_mb="$(read_profile_value "${profile}" carveout_target_mb)"
    ttm_action="$(read_profile_value "${profile}" ttm_action)"
    ttm_gb="$(read_profile_value "${profile}" ttm_gb)"
    min_reserve="$(read_profile_value "${profile}" min_os_reserve_gb)"

    [[ -n "${carveout_mb}" ]] || { echo "Missing carveout_target_mb for ${profile}"; return 1; }

    total_mb="$(total_ram_mb)"
    total_gb="$(total_ram_gb)"
    max_ttm_gb="$(python3 -c "print(${total_gb} * ${MAX_TTM_PERCENT} / 100)")"

    if [[ "${ttm_action}" == "clear" ]]; then
        ttm_gb="0"
    elif [[ -n "${ttm_gb}" ]]; then
        :
    else
        ttm_gb="0"
    fi

    local needed_mb
    needed_mb="$(python3 -c "
c = ${carveout_mb}
t = float('${ttm_gb}') * 1024 if ${ttm_gb} else 0
r = float('${min_reserve}') * 1024
print(int(c + t + r))
")"

    if (( needed_mb > total_mb )); then
        echo "Refusing: profile needs ~$(( needed_mb / 1024 )) GB but only ~${total_gb} GB RAM installed."
        echo "  carveout ${carveout_mb} MB + TTM ${ttm_gb} GB + OS reserve ${min_reserve} GB"
        return 1
    fi

    if [[ "${ttm_gb}" != "0" ]] && python3 -c "exit(0 if float('${ttm_gb}') > float('${max_ttm_gb}') else 1)"; then
        echo "Warning: TTM ${ttm_gb} GB exceeds ${MAX_TTM_PERCENT}% of RAM (~${max_ttm_gb} GB)."
    fi

    return 0
}

write_ttm_config() {
    local gb="$1"
    local pages
    pages="$(gb_to_pages "${gb}")"
    cat > "${TTM_CONF}" <<EOF
# Written by z13-memory-profiles on $(date -Iseconds)
options ttm pages_limit=${pages}
options ttm page_pool_size=${pages}
EOF
    log "Wrote ${TTM_CONF} (${gb} GB -> ${pages} pages)"
}

clear_ttm_config() {
    if [[ -f "${TTM_CONF}" ]]; then
        rm -f "${TTM_CONF}"
        log "Removed ${TTM_CONF}"
    else
        log "No ${TTM_CONF} to remove (already default)"
    fi
}

regenerate_initramfs() {
    if command -v dracut >/dev/null 2>&1; then
        log "Regenerating initramfs (dracut)..."
        dracut --force --regenerate-all 2>/dev/null || dracut --force
        return
    fi
    if command -v update-initramfs >/dev/null 2>&1; then
        update-initramfs -u -k all
        return
    fi
    if command -v mkinitcpio >/dev/null 2>&1; then
        mkinitcpio -P
        return
    fi
    log "No initramfs tool found; reboot may still apply modprobe.d on next boot."
}

set_carveout_index() {
    local index="$1"
    local uma_base
    uma_base="$(find_uma_base)"
    echo "${index}" > "${uma_base}/carveout"
    log "Set ${uma_base}/carveout -> ${index} (applies next boot)"
}

apply_profile() {
    local profile="$1"
    require_root

    local carveout_mb ttm_action ttm_gb label desc

    label="$(read_profile_value "${profile}" label)"
    desc="$(read_profile_value "${profile}" description)"
    carveout_mb="$(read_profile_value "${profile}" carveout_target_mb)"
    ttm_action="$(read_profile_value "${profile}" ttm_action)"
    ttm_gb="$(read_profile_value "${profile}" ttm_gb)"

    validate_profile "${profile}" || return 1

    log "Applying profile: ${label} (${profile})"
    log "${desc}"

    local uma_base carveout_line carveout_index
    if uma_base="$(find_uma_base)"; then
        mapfile -t carveout_pick < <(find_carveout_index "${uma_base}" "${carveout_mb}")
        carveout_index="${carveout_pick[0]}"
        carveout_line="${carveout_pick[1]:-}"
        set_carveout_index "${carveout_index}"
        log "Carveout: ${carveout_line}"
    else
        log "WARN: UMA sysfs missing — fixed VRAM unchanged."
        log "      Use Armoury Crate in Windows once, or update kernel/firmware."
    fi

    if [[ "${ttm_action}" == "clear" ]]; then
        clear_ttm_config
    elif [[ -n "${ttm_gb}" ]]; then
        write_ttm_config "${ttm_gb}"
    fi

    regenerate_initramfs

    cat <<EOF
Profile "${label}" staged successfully.

Changes take effect after reboot:
  • Fixed VRAM (carveout): target ${carveout_mb} MB
  • Dynamic GPU pool (TTM): $(
        if [[ "${ttm_action}" == "clear" ]]; then echo "kernel default (custom config removed)"; else echo "${ttm_gb} GB"; fi
    )

Verify after reboot:
  sudo dmesg | grep -i 'amdgpu.*memory'
  cat /sys/module/ttm/parameters/pages_limit 2>/dev/null || true
EOF
}

show_status() {
    local total_gb uma_line ttm_gb
    total_gb="$(total_ram_gb)"
    echo "=== Z13 Memory Status ==="
    echo "Total RAM: ~${total_gb} GB"
    echo ""

    if uma_base="$(find_uma_base)"; then
        echo "UMA carveout (fixed VRAM):"
        echo "  Path: ${uma_base}"
        echo "  Active index: $(cat "${uma_base}/carveout" 2>/dev/null || echo '?')"
        echo "  Options:"
        sed 's/^/    /' "${uma_base}/carveout_options"
    else
        echo "UMA carveout: not exposed on this kernel (use Windows Armoury Crate for fixed VRAM)"
    fi
    echo ""

    echo "TTM dynamic pool:"
    if [[ -r "${TTM_PARAM}" ]]; then
        ttm_gb="$(current_ttm_gb)"
        echo "  Active limit: ~${ttm_gb} GB"
    else
        echo "  ${TTM_PARAM} not readable"
    fi
    if [[ -f "${TTM_CONF}" ]]; then
        echo "  Pending config: ${TTM_CONF}"
        sed 's/^/    /' "${TTM_CONF}"
    else
        echo "  No custom ${TTM_CONF}"
    fi
    echo ""

    if command -v amd-ttm >/dev/null 2>&1; then
        echo "amd-ttm:"
        amd-ttm 2>/dev/null | sed 's/^/  /' || true
    fi
}