#!/bin/bash
# GPU Hot-Path Guard — detects CPU roundtrip leaks and GPU perf killers.
#
# Production HFT system on H100 GPUs. Every microsecond matters.
# This guard catches ALL patterns that cause GPU pipeline stalls,
# unnecessary CPU↔GPU memory transfers, or heap allocations in hot paths.
#
# Usage:
#   scripts/gpu-hotpath-guard.sh <file>          # check single file
#   scripts/gpu-hotpath-guard.sh --staged        # check all staged .rs files
#   scripts/gpu-hotpath-guard.sh --diff           # check unstaged changes only
#   scripts/gpu-hotpath-guard.sh --all            # check all hot-path .rs files
#
# Exit codes:
#   0  = clean (no leaks, or file not in hot path)
#   1  = GPU perf violation detected in hot-path file
#
# Suppress reviewed boundary points with:  // gpu-ok: <reason>
# NOTE: New code should NEVER use // gpu-ok: — eliminate the transfer instead.

set -euo pipefail

# ── Hot-path file patterns ──────────────────────────────────────────────
# Only files matching these patterns are checked. Everything else is ignored.
HOT_PATHS=(
    # DQN training hot loop (entire trainer directory)
    "trainers/dqn/"
    "cuda_pipeline/"
    "gpu_action_selector"
    "gpu_portfolio"
    "gpu_replay_buffer"
    "gpu_experience_collector"
    "gpu_backtest_evaluator"
    "gpu_training_guard"
    "gpu_weights"
    "mixed_precision.rs"
    # DQN crate — all modules (inference, networks, replay, regularization)
    "ml-dqn/src/"
    # PPO training + inference
    "trainers/ppo"
    "ml-ppo/src/"
    # Supervised model forward passes
    "ml-supervised/src/"
    # Hyperopt backtest inner loops
    "hyperopt/adapters/"
    "hyperopt/walk_forward.rs"
    # Evaluation binaries (GPU inference paths)
    "examples/evaluate_baseline.rs"
    "examples/train_baseline"
    # Ensemble inference
    "ensemble/"
    # Feature extraction (GPU-resident)
    "ml-features/src/"
    # Flash attention (GPU-only computation)
    "flash_attention/"
)

# ── GPU perf violation patterns ────────────────────────────────────────
# Each pattern is a grep -E regex. Lines with "// gpu-ok:" are excluded.
# Doc comments (/// or //!) and Storage match arms (error guards) are excluded.
LEAK_PATTERNS=(
    # ─── GPU→CPU data transfers (tensor → CPU memory) ───
    # Each of these forces a CUDA stream synchronize + PCIe DMA transfer.
    # On H100: ~5-15μs per transfer, blocks entire GPU pipeline.
    '\.to_vec1'
    '\.to_vec2'
    '\.to_vec3'
    '\.to_scalar'
    '\.to_device\(&Device::Cpu\)'
    '\.to_device\(&candle_core::Device::Cpu\)'
    '\.get\([0-9]+\)\?\.to_vec'

    # ─── CPU→GPU data transfers (CPU heap alloc + memcpy to VRAM) ───
    # Tensor::from_vec: allocates Vec on CPU heap, copies to GPU.
    # On H100: CPU alloc ~0.5-2μs + PCIe DMA ~5-15μs = 5-17μs per call.
    # Must be eliminated from inner loops. Acceptable at data-load boundaries.
    'Tensor::from_vec'
    'Tensor::from_slice'

    # ─── CPU device usage (creating anything on CPU in GPU code) ───
    '&Device::Cpu'

    # ─── Dead GPU transfers (underscore-hidden variables suppress clippy) ───
    # let _foo = Tensor::from_vec/from_slice — allocated on GPU, never used.
    # let _foo = tensor.to_vec1/to_scalar — forced GPU→CPU sync, result discarded.
    'let _\w+\s*=.*Tensor::from_vec'
    'let _\w+\s*=.*Tensor::from_slice'
    'let _\w+\s*=.*\.to_vec[123]'
    'let _\w+\s*=.*\.to_scalar'

    # NOTE: memcpy_dtoh/htod and .synchronize() are not checked here —
    # they are CUDA implementation primitives in cuda_pipeline/.
    # Audit those separately with: scripts/cuda-perf-audit.sh
)

# ── Excluded paths (test-only infra that deliberately uses CPU) ────────
EXCLUDE_PATHS=(
    "smoke_tests/"
    "_test.rs"
    "_tests.rs"
    "performance_tests.rs"
    "performance_validation.rs"
    "demo_dqn.rs"
)

is_hot_path() {
    local file="$1"
    # Check exclusions first
    for excl in "${EXCLUDE_PATHS[@]}"; do
        if [[ "$file" == *"$excl"* ]]; then
            return 1
        fi
    done
    for pattern in "${HOT_PATHS[@]}"; do
        if [[ "$file" == *"$pattern"* ]]; then
            return 0
        fi
    done
    return 1
}

check_file() {
    local file="$1"
    local found=0

    if ! is_hot_path "$file"; then
        return 0
    fi

    # Find the line number where #[cfg(test)] starts — everything after is test code
    local test_start
    test_start=$(grep -n '#\[cfg(test)\]' "$file" 2>/dev/null | head -1 | cut -d: -f1 || echo "999999")

    for pattern in "${LEAK_PATTERNS[@]}"; do
        # Find matches, exclude: suppressed lines, doc comments, Storage match arms, regular comments
        local hits
        hits=$(grep -nE "$pattern" "$file" 2>/dev/null \
            | grep -v '// gpu-ok:' \
            | grep -v '[0-9]*:\s*///' \
            | grep -v '[0-9]*:\s*//!' \
            | grep -v '[0-9]*:\s*//' \
            | grep -v 'Storage::Cpu' \
            | grep -v 'Storage::Metal' \
            || true)

        # Filter out lines in test code (line number >= test_start)
        if [ -n "$hits" ]; then
            local filtered=""
            while IFS= read -r line; do
                local lineno
                lineno=$(echo "$line" | cut -d: -f1)
                if [ "$lineno" -lt "$test_start" ] 2>/dev/null; then
                    filtered="${filtered}${line}"$'\n'
                fi
            done <<< "$hits"
            filtered="${filtered%$'\n'}"  # trim trailing newline

            if [ -n "$filtered" ]; then
                if [ "$found" -eq 0 ]; then
                    echo "GPU HOT-PATH LEAK: $file"
                    found=1
                fi
                echo "$filtered" | while IFS= read -r line; do
                    echo "  $line"
                done
            fi
        fi
    done

    return $found
}

# ── Entrypoint ──────────────────────────────────────────────────────────
FILES=()
MODE="file"

if [ "${1:-}" = "--staged" ]; then
    MODE="staged"
    while IFS= read -r f; do
        FILES+=("$f")
    done < <(git diff --cached --name-only --diff-filter=ACM | grep '\.rs$' || true)
elif [ "${1:-}" = "--diff" ]; then
    MODE="diff"
    while IFS= read -r f; do
        FILES+=("$f")
    done < <(git diff --name-only | grep '\.rs$' || true)
elif [ "${1:-}" = "--all" ]; then
    MODE="all"
    while IFS= read -r f; do
        FILES+=("$f")
    done < <(find crates -name '*.rs' -type f 2>/dev/null)
elif [ -n "${1:-}" ]; then
    FILES=("$1")
else
    echo "Usage: $0 <file|--staged|--diff|--all>"
    exit 0
fi

LEAKED=0
LEAK_COUNT=0
for file in "${FILES[@]}"; do
    if [ -f "$file" ]; then
        if ! check_file "$file"; then
            LEAKED=1
            LEAK_COUNT=$((LEAK_COUNT + 1))
        fi
    fi
done

if [ "$LEAKED" -eq 1 ]; then
    echo ""
    echo "⛔ $LEAK_COUNT file(s) with GPU perf violations in hot paths."
    echo "   MUST be eliminated — do not suppress with // gpu-ok:"
    echo "   Move transfers out of hot paths or keep data GPU-resident."
    exit 1
fi

exit 0