#!/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"
    # 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]'

    # ─── Dead CPU fallback paths ───
    # #[cfg(not(feature = "cuda"))] blocks are dead code — project is CUDA-only.
    # They add maintenance burden and hide the real code path.
    '#\[cfg\(not\(feature\s*=\s*"cuda"\)\)\]'

    # ─── CPU-side batch data mutation (GPU data wrongly processed on CPU) ───
    # Functions taking mutable Vec<f32> batch arrays indicate CPU-side processing
    # that should be a GPU kernel or Candle tensor ops.
    # Caught: CPU HER relabeling, CPU normalization, CPU reward shaping, etc.
    'mut states: Vec<f32>'
    'mut next_states: Vec<f32>'
    'mut rewards: Vec<f32>'

    # ─── Candle tensor wrapping in per-step hot path ───
    # Creating Candle Tensors from CudaSlice in the training loop adds
    # Rust-side allocation + dispatch overhead. Use raw CUDA kernels instead.
    # Tensor::new() creates a CPU scalar → GPU copy (HtoD for 4 bytes).
    # CudaStorage::wrap_cuda_slice takes ownership — blocks reuse patterns.
    'CudaStorage::wrap_cuda_slice'

    # 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
    #
    # NOTE: storage_and_layout() is NOT flagged — it's a zero-cost pointer
    # unwrap for extracting CudaSlice refs from existing Tensors.
)

# ── Excluded paths ────────────────────────────────────────────────────
EXCLUDE_PATHS=(
    # One-time data loading at startup, not per-step hot path
    "trainers/dqn/data_loading.rs"
    # Smoke tests (not training hot path)
    "trainers/dqn/smoke_tests/"
    # Legacy non-branching QNetwork — CPU API boundary, not production hot path
    "ml-dqn/src/network.rs"
    "ml-dqn/src/ensemble_network.rs"
    # Model inference boundaries — return CPU types by design (EnsemblePrediction)
    # Training hot path uses GPU-resident tensors directly
    "ml-supervised/src/tft/varmap_quantization.rs"
    # tensor_to_predictions() returns Vec<Price> — inference output boundary
    "ml-supervised/src/tft/hft_optimizations.rs"
    # evaluate_cpu() is slow fallback before precompute_gpu_grid(); GPU path exists
    "ml-supervised/src/kan/spline.rs"
    # SearchSorted/CumsumOp CustomOp2 impls must return CudaStorage (Candle API boundary)
    # Runs once per PER sampling batch, not inside per-step training kernel loop
    "ml-dqn/src/gpu_replay_buffer.rs"
)

# ── 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).
        # Tracks brace depth to skip entire functions/blocks, not just 2 lines.
        if [ -n "$hits" ]; then
            local skip_ranges
            skip_ranges=$(python3 -c "
import re, sys
lines = open(sys.argv[1]).readlines()
ranges = []
for i, line in enumerate(lines):
    if re.search(r'#\[cfg\(not\(feature.*cuda', line):
        start = i
        depth = 0
        found_open = False
        for j in range(i, len(lines)):
            for ch in lines[j]:
                if ch == '{':
                    depth += 1
                    found_open = True
                elif ch == '}':
                    depth -= 1
            if found_open and depth == 0:
                ranges.append((start + 1, j + 1))  # 1-indexed
                break
        else:
            # Single expression (no braces): skip attribute + next line
            ranges.append((start + 1, min(start + 2, len(lines))))
for s, e in ranges:
    print(f'{s} {e}')
" "$file" 2>/dev/null || true)
            if [ -n "$skip_ranges" ]; then
                hits=$(echo "$hits" | while IFS= read -r line; do
                    local lineno
                    lineno=$(echo "$line" | cut -d: -f1)
                    local skip=false
                    while IFS=' ' read -r range_start range_end; do
                        if [ "$lineno" -ge "$range_start" ] && [ "$lineno" -le "$range_end" ]; then
                            skip=true
                            break
                        fi
                    done <<< "$skip_ranges"
                    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