#!/bin/bash
# =============================================================================
# TEST SCRIPT: Optimized Dockerfile Validation
# =============================================================================
# Purpose: Validate Dockerfile.runpod.optimized meets all requirements
# Usage: ./scripts/test_optimized_dockerfile.sh
# =============================================================================

set -euo pipefail

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

# Test counters
TESTS_RUN=0
TESTS_PASSED=0
TESTS_FAILED=0

# Logging helpers
log_info() {
    echo -e "${GREEN}[INFO]${NC} $*"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $*"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $*"
}

log_test() {
    TESTS_RUN=$((TESTS_RUN + 1))
    echo -e "\n${YELLOW}[TEST $TESTS_RUN]${NC} $*"
}

log_pass() {
    TESTS_PASSED=$((TESTS_PASSED + 1))
    echo -e "${GREEN}  ✓ PASS${NC}: $*"
}

log_fail() {
    TESTS_FAILED=$((TESTS_FAILED + 1))
    echo -e "${RED}  ✗ FAIL${NC}: $*"
}

# Cleanup function
cleanup() {
    log_info "Cleaning up test artifacts..."
    docker rm -f foxhunt-test-container 2>/dev/null || true
    rm -rf /tmp/foxhunt-test-volume 2>/dev/null || true
}

trap cleanup EXIT

# =============================================================================
# TEST 1: Build Minimal Image (No SSH)
# =============================================================================
log_test "Build minimal image (no SSH)"

if docker build -f Dockerfile.runpod.optimized \
    -t jgrusewski/foxhunt:test-minimal \
    . > /tmp/build-minimal.log 2>&1; then
    log_pass "Minimal image built successfully"
else
    log_fail "Minimal image build failed"
    cat /tmp/build-minimal.log
    exit 1
fi

# =============================================================================
# TEST 2: Build Debug Image (SSH Enabled)
# =============================================================================
log_test "Build debug image (SSH enabled)"

if docker build -f Dockerfile.runpod.optimized \
    --build-arg INSTALL_SSH=true \
    -t jgrusewski/foxhunt:test-ssh \
    . > /tmp/build-ssh.log 2>&1; then
    log_pass "Debug image built successfully"
else
    log_fail "Debug image build failed"
    cat /tmp/build-ssh.log
    exit 1
fi

# =============================================================================
# TEST 3: Verify Image Sizes
# =============================================================================
log_test "Verify image sizes"

MINIMAL_SIZE=$(docker images jgrusewski/foxhunt:test-minimal --format "{{.Size}}")
SSH_SIZE=$(docker images jgrusewski/foxhunt:test-ssh --format "{{.Size}}")

log_info "Minimal image size: $MINIMAL_SIZE"
log_info "SSH image size: $SSH_SIZE"

# Convert sizes to MB for comparison (assumes GB/MB format)
MINIMAL_MB=$(echo "$MINIMAL_SIZE" | sed 's/GB/*1024/; s/MB//; s/KB\/1024/' | bc 2>/dev/null || echo "0")
SSH_MB=$(echo "$SSH_SIZE" | sed 's/GB/*1024/; s/MB//; s/KB\/1024/' | bc 2>/dev/null || echo "0")

if [ "$(echo "$MINIMAL_MB < 3500" | bc)" -eq 1 ]; then
    log_pass "Minimal image size acceptable: $MINIMAL_SIZE (target: <3.5GB)"
else
    log_fail "Minimal image too large: $MINIMAL_SIZE (target: <3.5GB)"
fi

if [ "$(echo "$SSH_MB < 4000" | bc)" -eq 1 ]; then
    log_pass "SSH image size acceptable: $SSH_SIZE (target: <4GB)"
else
    log_fail "SSH image too large: $SSH_SIZE (target: <4GB)"
fi

# =============================================================================
# TEST 4: Verify CUDA Runtime Base (Not Devel)
# =============================================================================
log_test "Verify using CUDA runtime base (not devel)"

if docker history jgrusewski/foxhunt:test-minimal | grep -q "13.0.0-runtime-ubuntu24.04"; then
    log_pass "Using CUDA runtime base image"
else
    log_fail "Not using CUDA runtime base image"
fi

# Check for absence of nvcc (should NOT be in runtime image)
if ! docker run --rm jgrusewski/foxhunt:test-minimal which nvcc 2>/dev/null; then
    log_pass "nvcc not present (correct for runtime image)"
else
    log_fail "nvcc present (indicates devel image used)"
fi

# =============================================================================
# TEST 5: Verify CUDA Libraries Present
# =============================================================================
log_test "Verify CUDA runtime libraries present"

CUDA_LIBS=(
    "/usr/local/cuda/lib64/libcuda.so.1"
    "/usr/local/cuda/lib64/libcurand.so.10"
    "/usr/local/cuda/lib64/libcublas.so.13"
    "/usr/local/cuda/lib64/libcublasLt.so.13"
)

for lib in "${CUDA_LIBS[@]}"; do
    if docker run --rm jgrusewski/foxhunt:test-minimal test -f "$lib"; then
        log_pass "Library present: $lib"
    else
        log_fail "Library missing: $lib"
    fi
done

# =============================================================================
# TEST 6: Verify cuDNN Present
# =============================================================================
log_test "Verify cuDNN runtime library present"

if docker run --rm jgrusewski/foxhunt:test-minimal \
    find /usr/lib -name "libcudnn.so*" 2>/dev/null | grep -q "libcudnn"; then
    log_pass "cuDNN runtime library present"
else
    log_fail "cuDNN runtime library missing"
fi

# =============================================================================
# TEST 7: Verify runpodctl Installed
# =============================================================================
log_test "Verify runpodctl CLI tool present"

if docker run --rm jgrusewski/foxhunt:test-minimal runpodctl --version 2>/dev/null; then
    log_pass "runpodctl installed and executable"
else
    log_fail "runpodctl missing or not executable"
fi

# =============================================================================
# TEST 8: Verify SSH Conditionally Installed
# =============================================================================
log_test "Verify SSH conditionally installed"

# Minimal image should NOT have SSH
if ! docker run --rm jgrusewski/foxhunt:test-minimal which sshd 2>/dev/null; then
    log_pass "SSH not present in minimal image (correct)"
else
    log_fail "SSH present in minimal image (should be disabled)"
fi

# SSH image SHOULD have SSH
if docker run --rm jgrusewski/foxhunt:test-ssh which sshd 2>/dev/null; then
    log_pass "SSH present in SSH image (correct)"
else
    log_fail "SSH missing in SSH image (should be enabled)"
fi

# =============================================================================
# TEST 9: Verify Entrypoint Scripts Present
# =============================================================================
log_test "Verify entrypoint scripts present and executable"

ENTRYPOINT_SCRIPTS=(
    "/entrypoint.sh"
    "/entrypoint-generic.sh"
)

for script in "${ENTRYPOINT_SCRIPTS[@]}"; do
    if docker run --rm jgrusewski/foxhunt:test-minimal test -x "$script"; then
        log_pass "Script executable: $script"
    else
        log_fail "Script missing or not executable: $script"
    fi
done

# =============================================================================
# TEST 10: Verify Volume Mount Validation
# =============================================================================
log_test "Verify entrypoint validates volume mount"

# Create test volume directory
mkdir -p /tmp/foxhunt-test-volume/binaries
mkdir -p /tmp/foxhunt-test-volume/test_data

# Test with volume mount (should succeed)
if docker run --rm \
    -v /tmp/foxhunt-test-volume:/runpod-volume \
    jgrusewski/foxhunt:test-minimal \
    /entrypoint-generic.sh --help 2>&1 | grep -q "Volume mount verified"; then
    log_pass "Entrypoint validates volume mount correctly"
else
    log_fail "Entrypoint volume validation failed"
fi

# Test without volume mount (should fail gracefully)
if docker run --rm \
    jgrusewski/foxhunt:test-minimal \
    /entrypoint-generic.sh --help 2>&1 | grep -q "ERROR.*not mounted"; then
    log_pass "Entrypoint fails gracefully without volume mount"
else
    log_fail "Entrypoint does not validate missing volume mount"
fi

# =============================================================================
# TEST 11: Verify Layer Count Reduced
# =============================================================================
log_test "Verify layer count reduced (vs current Dockerfile)"

MINIMAL_LAYERS=$(docker history jgrusewski/foxhunt:test-minimal --no-trunc | wc -l)
log_info "Optimized image layers: $MINIMAL_LAYERS"

# Current image has ~15 layers, target is <10
if [ "$MINIMAL_LAYERS" -lt 10 ]; then
    log_pass "Layer count acceptable: $MINIMAL_LAYERS (target: <10)"
else
    log_fail "Layer count too high: $MINIMAL_LAYERS (target: <10)"
fi

# =============================================================================
# TEST 12: Verify No Build Tools Present
# =============================================================================
log_test "Verify no build tools present (security)"

BUILD_TOOLS=("gcc" "g++" "make" "cmake" "git" "wget" "curl")

for tool in "${BUILD_TOOLS[@]}"; do
    if ! docker run --rm jgrusewski/foxhunt:test-minimal which "$tool" 2>/dev/null; then
        log_pass "Build tool NOT present: $tool (correct)"
    else
        log_fail "Build tool present: $tool (security risk)"
    fi
done

# =============================================================================
# TEST 13: Verify GPU Access (if GPU available)
# =============================================================================
log_test "Verify GPU access (if available)"

if command -v nvidia-smi &> /dev/null; then
    if docker run --rm --gpus all jgrusewski/foxhunt:test-minimal nvidia-smi > /dev/null 2>&1; then
        log_pass "GPU accessible from container"
    else
        log_warn "GPU access failed (may not be critical if no GPU present)"
    fi
else
    log_warn "nvidia-smi not found, skipping GPU test"
fi

# =============================================================================
# TEST 14: Compare Size vs Current Image
# =============================================================================
log_test "Compare size reduction vs current image"

if docker images jgrusewski/foxhunt:latest --format "{{.Size}}" | grep -q "GB"; then
    CURRENT_SIZE=$(docker images jgrusewski/foxhunt:latest --format "{{.Size}}")
    log_info "Current image size: $CURRENT_SIZE"
    log_info "Optimized image size: $MINIMAL_SIZE"

    # Calculate approximate savings
    CURRENT_GB=$(echo "$CURRENT_SIZE" | sed 's/GB//')
    MINIMAL_GB=$(echo "$MINIMAL_SIZE" | sed 's/GB//')
    SAVINGS=$(echo "scale=1; ($CURRENT_GB - $MINIMAL_GB) / $CURRENT_GB * 100" | bc 2>/dev/null || echo "0")

    log_pass "Size reduction: ~${SAVINGS}% (${CURRENT_SIZE} → ${MINIMAL_SIZE})"
else
    log_warn "Current image not found, skipping size comparison"
fi

# =============================================================================
# SUMMARY
# =============================================================================
echo ""
echo "======================================================================="
echo "TEST SUMMARY"
echo "======================================================================="
echo "Tests Run:    $TESTS_RUN"
echo "Tests Passed: $TESTS_PASSED"
echo "Tests Failed: $TESTS_FAILED"
echo ""

if [ "$TESTS_FAILED" -eq 0 ]; then
    echo -e "${GREEN}✓ ALL TESTS PASSED${NC}"
    echo ""
    echo "Next Steps:"
    echo "  1. Push to Docker Hub:"
    echo "     docker push jgrusewski/foxhunt:test-minimal"
    echo ""
    echo "  2. Deploy test pod on Runpod:"
    echo "     ./scripts/runpod_deploy_production.py --image jgrusewski/foxhunt:test-minimal"
    echo ""
    echo "  3. Monitor startup time (target: <2 min)"
    echo ""
    exit 0
else
    echo -e "${RED}✗ TESTS FAILED: $TESTS_FAILED${NC}"
    echo ""
    echo "Review failures above and fix Dockerfile.runpod.optimized"
    exit 1
fi