#!/usr/bin/env bash
# ptx-cache-invalidate.sh — Invalidate PTX cache when CUDA kernels change.
#
# The Rust runtime caches compiled PTX in $CARGO_TARGET_DIR/.ptx_cache/ (or
# /tmp/.ptx_cache/).  The cache key includes a SHA-256 of the kernel source,
# so changed .cu files naturally miss.  However, if the CI PVC preserves the
# cache across runs and the build binary is stale, old PTX can survive.
#
# This script computes a content hash of all .cu/.cuh kernel source files and
# compares it against a stamp file in the cache directory.  If the hash differs
# (any kernel source changed), the entire PTX cache is purged, forcing nvcc
# recompilation on the next run.
#
# Usage:  scripts/ptx-cache-invalidate.sh [cache_dir]
#   cache_dir  defaults to $CARGO_TARGET_DIR/.ptx_cache, then /tmp/.ptx_cache

set -euo pipefail

KERNEL_DIRS="crates/ml/src/cuda_pipeline crates/ml-dqn/src crates/ml-core/src/cuda_autograd crates/ml-ppo/src/cuda_nn"
CACHE_DIR="${1:-${CARGO_TARGET_DIR:+${CARGO_TARGET_DIR}/.ptx_cache}}"
CACHE_DIR="${CACHE_DIR:-/tmp/.ptx_cache}"

# Compute content hash of CUDA sources AND build.rs (nvcc flags).
# build.rs changes (e.g. removing --use_fast_math) produce different cubins
# even when .cu source is unchanged — must invalidate cache.
HASH=$( {
    find $KERNEL_DIRS -type f \( -name '*.cu' -o -name '*.cuh' \) -exec sha256sum {} + 2>/dev/null
    find crates/ml crates/ml-dqn crates/ml-core crates/ml-ppo -maxdepth 1 -name 'build.rs' -exec sha256sum {} + 2>/dev/null
} | sort | sha256sum | cut -d' ' -f1)

STAMP_FILE="${CACHE_DIR}/.kernel_hash"

if [ -f "$STAMP_FILE" ]; then
    OLD_HASH=$(cat "$STAMP_FILE")
    if [ "$HASH" = "$OLD_HASH" ]; then
        echo "PTX cache valid — kernel hash unchanged ($HASH)"
        exit 0
    fi
    echo "PTX cache STALE — kernel hash changed"
    echo "  old: $OLD_HASH"
    echo "  new: $HASH"
else
    echo "PTX cache stamp missing — first run or cache cleared"
fi

# Purge stale PTX cache
if [ -d "$CACHE_DIR" ]; then
    PTX_COUNT=$(find "$CACHE_DIR" -name '*.ptx' -type f 2>/dev/null | wc -l)
    echo "Purging ${PTX_COUNT} cached PTX files from ${CACHE_DIR}"
    rm -f "${CACHE_DIR}"/*.ptx
fi

# Write new stamp
mkdir -p "$CACHE_DIR"
echo "$HASH" > "$STAMP_FILE"
echo "PTX cache stamp updated: $HASH"