#!/bin/bash
# ============================================================================
# NVCC Wrapper for Air.rs — v1.1.5
#
# Fixes applied transparently to every nvcc invocation:
#   1. -fPIC         — required for shared-library (.so) builds on Linux
#   2. -arch=sm_XX   — compile for the *actual* installed GPU ISA, not the
#                      ancient sm_52 default that NVCC uses without this flag.
#                      Detected at build-time via nvidia-smi.
#   3. -O3           — maximum compiler optimization
#   4. --use_fast_math — hardware-accelerated transcendentals (exp, tanh, …)
# ============================================================================

set -euo pipefail

ARGS=()
HAS_PIC=0
HAS_ARCH=0
HAS_OPT=0
HAS_FAST_MATH=0

for arg in "$@"; do
    ARGS+=("$arg")
    case "$arg" in
        -fPIC|-fpic) HAS_PIC=1 ;;
        -arch=*|-gencode*) HAS_ARCH=1 ;;
        -O*) HAS_OPT=1 ;;
        --use_fast_math) HAS_FAST_MATH=1 ;;
    esac
done

# ── Find the real nvcc (avoid recursion into this wrapper) ────────────────────
REAL_NVCC=""
while IFS= read -r candidate; do
    # Skip this wrapper
    case "$candidate" in
        */Air.rs/scripts/nvcc) continue ;;
    esac
    REAL_NVCC="$candidate"
    break
done < <(which -a nvcc 2>/dev/null || true)

# Fallback: check common CUDA install paths
if [[ -z "$REAL_NVCC" ]]; then
    for p in \
        "${CUDA_HOME:-}/bin/nvcc" \
        "/usr/local/cuda/bin/nvcc" \
        "/usr/bin/nvcc"
    do
        [[ -x "$p" ]] && { REAL_NVCC="$p"; break; }
    done
fi

[[ -z "$REAL_NVCC" ]] && { echo "nvcc wrapper: real nvcc not found" >&2; exit 1; }

# ── 1. Position Independent Code ─────────────────────────────────────────────
if [[ "$OSTYPE" == "linux-gnu"* && $HAS_PIC -eq 0 ]]; then
    ARGS+=("-Xcompiler" "-fPIC")
fi

# ── 2. GPU Architecture Targeting ────────────────────────────────────────────
# Only inject -arch if the caller hasn't already set one.
# This is the key optimization: compiling for sm_89 (Ada Lovelace), sm_90
# (Hopper), or sm_100 (Blackwell) instead of the NVCC default (sm_52) gives
# access to newer ISA features, better register alloc, and HW-specific paths.
if [[ $HAS_ARCH -eq 0 ]]; then
    # 1. Detect NVCC version limits
    CUDA_VER=$("$REAL_NVCC" --version 2>&1 | grep "release" | sed 's/.*release //' | sed 's/,.*//' || echo "")
    CUDA_MAJOR=0
    CUDA_MINOR=0
    if [[ "$CUDA_VER" =~ ^([0-9]+)\.([0-9]+) ]]; then
        CUDA_MAJOR="${BASH_REMATCH[1]}"
        CUDA_MINOR="${BASH_REMATCH[2]}"
    fi

    # Determine maximum compiler capability
    MAX_CAPABILITY=75
    if (( CUDA_MAJOR > 12 )) || { [[ "$CUDA_MAJOR" -eq 12 ]] && (( CUDA_MINOR >= 8 )); }; then
        MAX_CAPABILITY=100
    elif [[ "$CUDA_MAJOR" -eq 12 ]]; then
        MAX_CAPABILITY=90
    elif [[ "$CUDA_MAJOR" -eq 11 ]] && (( CUDA_MINOR >= 8 )); then
        MAX_CAPABILITY=90
    elif [[ "$CUDA_MAJOR" -eq 11 ]] && (( CUDA_MINOR >= 1 )); then
        MAX_CAPABILITY=86
    elif [[ "$CUDA_MAJOR" -eq 11 ]]; then
        MAX_CAPABILITY=80
    fi

    # 2. Get targets representation
    TARGETS=""
    if [[ -n "${NVCC_ARCH:-}" ]]; then
        TARGETS="${NVCC_ARCH}"
    else
        # Auto-detect all host GPUs
        if DETECTED_CAPS=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | tr -d '.' | paste -sd, -); then
            TARGETS="$DETECTED_CAPS"
        fi
    fi

    # 3. Parse targets into capabilities
    ARCH_STRS=()
    if [[ "$TARGETS" == "all" ]]; then
        ALL_ARCHS=(70 75 80 86 89 90 100)
        for cur in "${ALL_ARCHS[@]}"; do
            if (( cur <= MAX_CAPABILITY )); then
                ARCH_STRS+=("$cur")
            fi
        done
    elif [[ -n "$TARGETS" ]]; then
        IFS=',' read -ra ADDR <<< "${TARGETS// /}"
        for arch in "${ADDR[@]}"; do
            digits=$(echo "$arch" | tr -cd '0-9')
            if [[ -n "$digits" ]]; then
                if (( digits <= MAX_CAPABILITY )); then
                    ARCH_STRS+=("$digits")
                fi
            fi
        done
    fi

    # Fallback to compiler's max capability if list resolved empty
    if [ ${#ARCH_STRS[@]} -eq 0 ]; then
        ARCH_STRS+=("$MAX_CAPABILITY")
    fi

    # 4. Generate -gencode arguments
    HIGHEST_ARCH=0
    for arch in "${ARCH_STRS[@]}"; do
        ARGS+=("-gencode" "arch=compute_${arch},code=sm_${arch}")
        if (( arch > HIGHEST_ARCH )); then
            HIGHEST_ARCH="$arch"
        fi
    done
    # Add virtual assembly target for forward compatibility (highest capability only)
    if (( HIGHEST_ARCH > 0 )); then
        ARGS+=("-gencode" "arch=compute_${HIGHEST_ARCH},code=compute_${HIGHEST_ARCH}")
    fi
fi

# ── 3. Optimization flags ─────────────────────────────────────────────────────
[[ $HAS_OPT       -eq 0 ]] && ARGS+=("-O3")
[[ $HAS_FAST_MATH -eq 0 ]] && ARGS+=("--use_fast_math")

exec "$REAL_NVCC" "${ARGS[@]}"
