#!/bin/bash
# =============================================================================
# Binary Sync Validation Test Script
# =============================================================================
# Purpose: Validate the binary volume sync fix works correctly
# Tests:
#   1. Local binary has correct timestamp metadata
#   2. S3 upload includes timestamp in metadata
#   3. S3 binary can be downloaded and verified
#   4. Timestamps are preserved and comparable
#
# Usage: ./test_binary_sync.sh <binary_name>
# Example: ./test_binary_sync.sh hyperopt_mamba2_demo
# =============================================================================

set -e

BINARY_NAME="$1"
if [ -z "$BINARY_NAME" ]; then
    echo "Usage: $0 <binary_name>"
    echo "Example: $0 hyperopt_mamba2_demo"
    exit 1
fi

# Configuration
LOCAL_PATH="target/release/examples/${BINARY_NAME}"
S3_KEY_CURRENT="binaries/current/${BINARY_NAME}"
S3_ENDPOINT="https://s3api-eur-is-1.runpod.io"
S3_BUCKET="se3zdnb5o4"
S3_PROFILE="runpod"

# Note: Upload creates two S3 objects:
#   1. binaries/timestamped/${BINARY_NAME}_YYYYMMDD_HHMMSS (timestamped version)
#   2. binaries/current/${BINARY_NAME} (symlink/copy to latest)
# This test validates the 'current' version which is what deployments use

echo "========================================================================"
echo "BINARY SYNC VALIDATION TEST"
echo "========================================================================"
echo "Binary: $BINARY_NAME"
echo "Local:  $LOCAL_PATH"
echo "S3:     s3://$S3_BUCKET/$S3_KEY_CURRENT"
echo "========================================================================"

# =============================================================================
# TEST 1: Verify local binary exists and has timestamp
# =============================================================================
echo ""
echo "TEST 1: Local Binary Validation"
echo "------------------------------------------------------------------------"

if [ ! -f "$LOCAL_PATH" ]; then
    echo "❌ FAIL: Local binary not found at $LOCAL_PATH"
    echo "   Build it with: cargo build --release --example $BINARY_NAME"
    exit 1
fi
echo "✅ Local binary exists"

# Get local file modification time
LOCAL_MTIME=$(stat -c %Y "$LOCAL_PATH" 2>/dev/null || stat -f %m "$LOCAL_PATH" 2>/dev/null)
LOCAL_SIZE=$(stat -c %s "$LOCAL_PATH" 2>/dev/null || stat -f %z "$LOCAL_PATH" 2>/dev/null)
LOCAL_TIMESTAMP=$(date -Iseconds -d @"$LOCAL_MTIME" 2>/dev/null || date -Iseconds -r "$LOCAL_MTIME" 2>/dev/null)

echo "✅ Local metadata extracted:"
echo "   Size:      $LOCAL_SIZE bytes ($(echo "scale=1; $LOCAL_SIZE/1024/1024" | bc) MB)"
echo "   Modified:  $LOCAL_TIMESTAMP"

# Calculate local SHA256
LOCAL_SHA256=$(sha256sum "$LOCAL_PATH" | awk '{print $1}')
echo "   SHA256:    $LOCAL_SHA256"

# =============================================================================
# TEST 2: Check S3 binary metadata (if exists)
# =============================================================================
echo ""
echo "TEST 2: S3 Binary Metadata Validation"
echo "------------------------------------------------------------------------"

# Check if S3 binary exists
if ! aws s3api head-object \
    --bucket "$S3_BUCKET" \
    --key "$S3_KEY_CURRENT" \
    --endpoint-url "$S3_ENDPOINT" \
    --profile "$S3_PROFILE" &>/dev/null; then
    echo "⚠️  S3 binary does not exist yet"
    echo "   Run: python3 scripts/runpod_deploy.py --skip-upload=false --dry-run"
    echo "   This will upload the binary with timestamp metadata"
    exit 0
fi

echo "✅ S3 binary exists"

# Get S3 metadata
S3_METADATA=$(aws s3api head-object \
    --bucket "$S3_BUCKET" \
    --key "$S3_KEY_CURRENT" \
    --endpoint-url "$S3_ENDPOINT" \
    --profile "$S3_PROFILE")

S3_SIZE=$(echo "$S3_METADATA" | jq -r '.ContentLength')
S3_MODIFIED=$(echo "$S3_METADATA" | jq -r '.LastModified')
S3_SHA256=$(echo "$S3_METADATA" | jq -r '.Metadata.sha256 // "N/A"')
S3_BUILD_TIMESTAMP=$(echo "$S3_METADATA" | jq -r '.Metadata.build_timestamp // "N/A"')
S3_UPLOADED=$(echo "$S3_METADATA" | jq -r '.Metadata.uploaded // "N/A"')

echo "✅ S3 metadata extracted:"
echo "   Size:            $S3_SIZE bytes ($(echo "scale=1; $S3_SIZE/1024/1024" | bc) MB)"
echo "   Last Modified:   $S3_MODIFIED"
echo "   SHA256:          $S3_SHA256"
echo "   Build Timestamp: $S3_BUILD_TIMESTAMP"
echo "   Uploaded:        $S3_UPLOADED"

# =============================================================================
# TEST 3: Validate timestamp metadata
# =============================================================================
echo ""
echo "TEST 3: Timestamp Validation"
echo "------------------------------------------------------------------------"

if [ "$S3_BUILD_TIMESTAMP" = "N/A" ]; then
    echo "⚠️  WARNING: S3 binary missing build_timestamp metadata"
    echo "   This binary was uploaded with old version of runpod_deploy.py"
    echo "   Re-upload to add timestamp: python3 scripts/runpod_deploy.py --force-upload --dry-run"
    exit 0
fi

echo "✅ S3 binary has timestamp metadata"

# Compare timestamps (allowing for timezone differences)
if [ "$LOCAL_TIMESTAMP" = "$S3_BUILD_TIMESTAMP" ]; then
    echo "✅ PASS: Timestamps match exactly"
    echo "   Local: $LOCAL_TIMESTAMP"
    echo "   S3:    $S3_BUILD_TIMESTAMP"
else
    echo "⚠️  Timestamps differ (expected if binary was rebuilt):"
    echo "   Local: $LOCAL_TIMESTAMP"
    echo "   S3:    $S3_BUILD_TIMESTAMP"

    # Check if local is newer (requires re-upload)
    LOCAL_EPOCH=$(date -d "$LOCAL_TIMESTAMP" +%s 2>/dev/null || date -j -f "%Y-%m-%dT%H:%M:%S" "$LOCAL_TIMESTAMP" +%s 2>/dev/null)
    S3_EPOCH=$(date -d "$S3_BUILD_TIMESTAMP" +%s 2>/dev/null || date -j -f "%Y-%m-%dT%H:%M:%S" "$S3_BUILD_TIMESTAMP" +%s 2>/dev/null)

    if [ "$LOCAL_EPOCH" -gt "$S3_EPOCH" ]; then
        echo "   ⚠️  Local binary is NEWER - re-upload required!"
        echo "      Run: python3 scripts/runpod_deploy.py --force-upload --dry-run"
    else
        echo "   ℹ️  S3 binary is same or newer - deployment will use S3 version"
    fi
fi

# =============================================================================
# TEST 4: Validate SHA256 checksum
# =============================================================================
echo ""
echo "TEST 4: Checksum Validation"
echo "------------------------------------------------------------------------"

if [ "$S3_SHA256" = "N/A" ]; then
    echo "⚠️  WARNING: S3 binary missing SHA256 metadata"
    echo "   Re-upload to add checksum: python3 scripts/runpod_deploy.py --force-upload --dry-run"
    exit 0
fi

echo "✅ S3 binary has SHA256 metadata"

# Compare checksums
if [ "$LOCAL_SHA256" = "$S3_SHA256" ]; then
    echo "✅ PASS: Checksums match - binaries are identical"
    echo "   SHA256: $LOCAL_SHA256"
else
    echo "❌ FAIL: Checksums DO NOT MATCH"
    echo "   Local:  $LOCAL_SHA256"
    echo "   S3:     $S3_SHA256"
    echo ""
    echo "   REQUIRED ACTION:"
    echo "   1. Force re-upload: python3 scripts/runpod_deploy.py --force-upload --dry-run"
    echo "   2. Or delete and re-upload:"
    echo "      aws s3 rm s3://$S3_BUCKET/$S3_KEY_CURRENT --endpoint-url $S3_ENDPOINT --profile $S3_PROFILE"
    echo "      python3 scripts/runpod_deploy.py --dry-run"
    exit 1
fi

# =============================================================================
# TEST 5: Validate binary arguments (for hyperopt binaries)
# =============================================================================
echo ""
echo "TEST 5: Binary Arguments Validation"
echo "------------------------------------------------------------------------"

if [[ "$BINARY_NAME" == hyperopt* ]]; then
    if ! "$LOCAL_PATH" --help 2>&1 | grep -q -- "--base-dir"; then
        echo "❌ FAIL: Local binary missing --base-dir argument"
        echo "   This binary was built BEFORE VarMap fix!"
        echo "   Rebuild: cargo clean && cargo build --release --example $BINARY_NAME"
        exit 1
    fi
    echo "✅ PASS: Local binary has --base-dir argument (VarMap fix applied)"

    # Test S3 binary (download temporarily)
    TEMP_S3="/tmp/${BINARY_NAME}_s3_test_$$"
    echo "   Downloading S3 binary for argument test..."
    aws s3 cp \
        "s3://${S3_BUCKET}/${S3_KEY_CURRENT}" \
        "$TEMP_S3" \
        --endpoint-url "$S3_ENDPOINT" \
        --profile "$S3_PROFILE" \
        --quiet

    chmod +x "$TEMP_S3"

    if ! "$TEMP_S3" --help 2>&1 | grep -q -- "--base-dir"; then
        echo "❌ FAIL: S3 binary missing --base-dir argument"
        echo "   This binary was uploaded BEFORE VarMap fix!"
        echo "   Re-upload: python3 scripts/runpod_deploy.py --force-upload --dry-run"
        rm -f "$TEMP_S3"
        exit 1
    fi
    echo "✅ PASS: S3 binary has --base-dir argument (VarMap fix applied)"

    rm -f "$TEMP_S3"
else
    echo "ℹ️  SKIP: Not a hyperopt binary, no argument validation needed"
fi

# =============================================================================
# SUMMARY
# =============================================================================
echo ""
echo "========================================================================"
echo "✅ ALL TESTS PASSED"
echo "========================================================================"
echo "Binary:           $BINARY_NAME"
echo "Local Size:       $(echo "scale=1; $LOCAL_SIZE/1024/1024" | bc) MB"
echo "Local Timestamp:  $LOCAL_TIMESTAMP"
echo "S3 Timestamp:     $S3_BUILD_TIMESTAMP"
echo "Checksum:         $LOCAL_SHA256"
echo ""
echo "DEPLOYMENT STATUS: ✅ READY"
echo "This binary is safe to deploy to Runpod."
echo "The entrypoint will sync it from S3 to volume on pod startup."
echo "========================================================================"