#!/usr/bin/env bash
# =============================================================================
# Training Infrastructure Consistency Tests
# =============================================================================
# Validates that Dockerfile.training, train.sh, build_docker_images.sh, and
# job-template.yaml all agree on binary names and model mappings.
#
# Run: bash infra/scripts/tests/test_training_infra_consistency.sh
# =============================================================================

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"

PASS=0
FAIL=0

pass() { echo "  PASS: $1"; PASS=$((PASS + 1)); }
fail() { echo "  FAIL: $1"; FAIL=$((FAIL + 1)); }

echo "=== Training Infrastructure Consistency Tests ==="
echo ""

# --- Test 1: Dockerfile.training builds all expected binaries ----------------
echo "[1] Dockerfile.training binary list"

DOCKERFILE="$REPO_ROOT/infra/docker/Dockerfile.training"
EXPECTED_BINS=(train_baseline_rl train_baseline_supervised evaluate_baseline hyperopt_baseline_rl hyperopt_baseline_supervised)

for bin in "${EXPECTED_BINS[@]}"; do
    if grep -q "$bin" "$DOCKERFILE"; then
        pass "Dockerfile builds $bin"
    else
        fail "Dockerfile missing $bin"
    fi
done

# --- Test 2: build_docker_images.sh EXPECTED_BINARIES matches ----------------
echo ""
echo "[2] build_docker_images.sh EXPECTED_BINARIES"

BUILD_SCRIPT="$REPO_ROOT/scripts/build_docker_images.sh"

# Only check the hyperopt binaries (that's what EXPECTED_BINARIES validates)
for bin in hyperopt_baseline_rl hyperopt_baseline_supervised; do
    if grep -q "\"$bin\"" "$BUILD_SCRIPT"; then
        pass "build_docker_images.sh validates $bin"
    else
        fail "build_docker_images.sh missing $bin from EXPECTED_BINARIES"
    fi
done

# --- Test 3: train.sh MODEL_BINARY covers all 10 models ---------------------
echo ""
echo "[3] train.sh MODEL_BINARY mapping completeness"

TRAIN_SCRIPT="$REPO_ROOT/infra/scripts/train.sh"
ALL_MODELS=(dqn ppo tft mamba2 tggn tlob liquid kan xlstm diffusion)

for model in "${ALL_MODELS[@]}"; do
    if grep -q "\\[$model\\]=" "$TRAIN_SCRIPT"; then
        pass "train.sh maps model=$model"
    else
        fail "train.sh missing mapping for model=$model"
    fi
done

# --- Test 4: train.sh maps RL models to train_baseline_rl -------------------
echo ""
echo "[4] train.sh RL model mapping"

RL_MODELS=(dqn ppo)
for model in "${RL_MODELS[@]}"; do
    if grep "\\[$model\\]=" "$TRAIN_SCRIPT" | grep -q "train_baseline_rl"; then
        pass "$model maps to train_baseline_rl"
    else
        fail "$model should map to train_baseline_rl"
    fi
done

# --- Test 5: train.sh maps supervised models to train_baseline_supervised ----
echo ""
echo "[5] train.sh supervised model mapping"

SUPERVISED_MODELS=(tft mamba2 tggn tlob liquid kan xlstm diffusion)
for model in "${SUPERVISED_MODELS[@]}"; do
    if grep "\\[$model\\]=" "$TRAIN_SCRIPT" | grep -q "train_baseline_supervised"; then
        pass "$model maps to train_baseline_supervised"
    else
        fail "$model should map to train_baseline_supervised"
    fi
done

# --- Test 6: Dockerfile COPY and strip loop are consistent -------------------
echo ""
echo "[6] Dockerfile build-and-copy consistency"

# Check the for loop that copies binaries matches the cargo build --example list
CARGO_EXAMPLES=$(grep -oP '(?<=--example )\S+' "$DOCKERFILE" | sort)
COPY_BINS=$(grep -oP '(?<=for bin in )\S.*?(?=;)' "$DOCKERFILE" | tr ' ' '\n' | sort)

if [ "$CARGO_EXAMPLES" = "$COPY_BINS" ]; then
    pass "Dockerfile cargo build examples match copy loop"
else
    fail "Dockerfile cargo build examples don't match copy loop"
    echo "    Build: $CARGO_EXAMPLES"
    echo "    Copy:  $COPY_BINS"
fi

# --- Test 7: job-template.yaml references valid binary -----------------------
echo ""
echo "[7] job-template.yaml binary reference"

JOB_TEMPLATE="$REPO_ROOT/infra/k8s/training/job-template.yaml"
if grep -q "train_baseline_supervised\|train_baseline_rl" "$JOB_TEMPLATE"; then
    pass "job-template.yaml references unified binary"
else
    fail "job-template.yaml doesn't reference unified binaries"
fi

# --- Test 8: No references to deleted binaries in infra ----------------------
echo ""
echo "[8] No stale binary references in infra/"

DELETED_BINS=(train_dqn train_tft train_ppo_parquet train_tft_parquet train_mamba2_parquet
              train_dqn_production train_rainbow train_ppo_extended train_continuous_ppo_parquet
              hyperopt_dqn_demo hyperopt_ppo_demo hyperopt_tft_demo hyperopt_mamba2_demo
              hyperopt_continuous_ppo_demo hyperopt_baseline)

STALE_FOUND=0
for bin in "${DELETED_BINS[@]}"; do
    # Search infra/ and scripts/ (excluding this test and docs/)
    if grep -rq --include='*.sh' --include='*.yaml' --include='*.yml' "$bin" \
        "$REPO_ROOT/infra/" "$REPO_ROOT/scripts/" 2>/dev/null | grep -v "test_training_infra" | grep -v "docs/" | head -1; then
        fail "Stale reference to deleted binary: $bin"
        STALE_FOUND=$((STALE_FOUND + 1))
    fi
done
if [ "$STALE_FOUND" -eq 0 ]; then
    pass "No stale binary references in infra/ or scripts/"
fi

# --- Test 9: Cargo.toml has example entries for all unified binaries ---------
echo ""
echo "[9] Cargo.toml example entries"

CARGO_TOML="$REPO_ROOT/crates/ml/Cargo.toml"
for bin in train_baseline_rl train_baseline_supervised hyperopt_baseline_rl hyperopt_baseline_supervised evaluate_baseline; do
    if grep -q "name = \"$bin\"" "$CARGO_TOML"; then
        pass "Cargo.toml has example entry: $bin"
    else
        fail "Cargo.toml missing example entry: $bin"
    fi
done

# --- Summary ------------------------------------------------------------------
echo ""
echo "=== Results: $PASS passed, $FAIL failed ==="

if [ "$FAIL" -gt 0 ]; then
    exit 1
fi