#!/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 bulk transfers (tensor array → CPU memory) ───
    # Each of these forces a CUDA stream synchronize + PCIe DMA transfer.
    # On H100: ~5-15μs per transfer, blocks entire GPU pipeline.
    # NOTE: .to_scalar() is NOT flagged — it reads a single value (4-8 bytes),
    # equivalent to the "final metric readback at epoch boundary" pattern.
    '\.to_vec1'
    '\.to_vec2'
    '\.to_vec3'
    '\.to_device\(&Device::Cpu\)'
    '\.to_device\(&candle_core::Device::Cpu\)'
    '\.get\([0-9]+\)\?\.to_vec'

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

    # ─── Dead GPU transfers (underscore-hidden variables suppress clippy) ───
    # let _foo = tensor.to_vec1 — forced GPU→CPU sync, result discarded.
    'let _\w+\s*=.*\.to_vec[123]'

    # NOTE: Tensor::from_vec/from_slice are NOT flagged — they are CPU→GPU
    # uploads (correct direction). Inner-loop abuse is caught by code review.
    #
    # 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 ────────────────────────────────────────────────────
EXCLUDE_PATHS=()

# ── Test module exclusion ─────────────────────────────────────────────
# Unit tests (#[cfg(test)] modules) are excluded — assertions inherently
# need scalar readbacks. Use .to_scalar() in tests (not flagged).
# Integration tests (crates/ml/tests/) ARE checked via HOT_PATHS.
find_test_module_start() {
    local file="$1"
    grep -n '^\s*#\[cfg(test)\]' "$file" 2>/dev/null | head -1 | cut -d: -f1
}

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 where #[cfg(test)] starts — skip all lines at or after that point
    local test_start
    test_start=$(find_test_module_start "$file")

    for pattern in "${LEAK_PATTERNS[@]}"; do
        # Find matches, exclude: doc comments, Storage match arms, regular comments
        local hits
        hits=$(grep -nE "$pattern" "$file" 2>/dev/null \
            | 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 inside #[cfg(test)] module
        if [ -n "$test_start" ] && [ -n "$hits" ]; then
            hits=$(echo "$hits" | while IFS= read -r line; do
                local lineno
                lineno=$(echo "$line" | cut -d: -f1)
                if [ "$lineno" -lt "$test_start" ]; then
                    echo "$line"
                fi
            done)
        fi

        # Filter out #[cfg(not(feature = "cuda"))] guarded code (dead code with CUDA).
        # The cfg attribute is on one line, the guarded expression on the next.
        if [ -n "$hits" ]; then
            local cfg_lines
            cfg_lines=$(grep -n '#\[cfg(not(feature.*cuda' "$file" 2>/dev/null | cut -d: -f1 || true)
            if [ -n "$cfg_lines" ]; then
                local skip_lines=""
                for cl in $cfg_lines; do
                    # Skip the cfg line itself AND the next line (the guarded expression)
                    skip_lines="$skip_lines $cl $((cl + 1))"
                done
                hits=$(echo "$hits" | while IFS= read -r line; do
                    local lineno
                    lineno=$(echo "$line" | cut -d: -f1)
                    local skip=false
                    for sl in $skip_lines; do
                        if [ "$lineno" = "$sl" ]; then
                            skip=true
                            break
                        fi
                    done
                    if [ "$skip" = false ]; then
                        echo "$line"
                    fi
                done)
            fi
        fi

        if [ -n "$hits" ]; then
            if [ "$found" -eq 0 ]; then
                echo "GPU HOT-PATH LEAK: $file"
                found=1
            fi
            echo "$hits" | while IFS= read -r line; do
                echo "  $line"
            done
        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