#!/bin/bash
# Test 03: Resource Enforcement Verification
# Verifies that resource limits are actually enforced using docker stats
# Usage: bash labs/lab-03-compute/tests/03-enforcement-test.sh

set -euo pipefail

# Color definitions
RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[1;33m'
BOLD='\033[1m'
NC='\033[0m'

# Get script directory
TEST_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LAB_DIR="$(cd "$TEST_DIR/.." && pwd)"

# Counter helpers
pass_count=0
fail_count=0

inc_pass() { ((pass_count++)) || true; }
inc_fail() { ((fail_count++)) || true; }

# Helper functions
print_header() {
    echo -e "${BLUE}╔═══════════════════════════════════════════════════════════════╗${NC}"
    echo -e "${BLUE}║${NC} ${BOLD}$1${NC}"
    echo -e "${BLUE}╚═══════════════════════════════════════════════════════════════╝${NC}"
}

print_test() {
    echo -e "\n${BLUE}[TEST]${NC} $1"
}

print_pass() {
    echo -e "  ${GREEN}[✓]${NC} $1"
    inc_pass
}

print_fail() {
    echo -e "  ${RED}[✗]${NC} $1"
    inc_fail
}

print_info() {
    echo -e "  ${BLUE}[i]${NC} $1"
}

print_warning() {
    echo -e "  ${YELLOW}[!]${NC} $1"
}

# Cleanup function
cleanup() {
    print_info "Cleaning up test containers..."
    cd "$LAB_DIR"
    docker compose down -v 2>/dev/null || true
}

# Set trap for cleanup
trap cleanup EXIT

# Main verification
print_header "Lab 03 Resource Enforcement Verification"

cd "$LAB_DIR"

# Test 1: Verify docker-compose.yml exists and is valid
print_test "Verifying docker-compose.yml exists and is valid"
if [[ ! -f "docker-compose.yml" ]]; then
    print_fail "docker-compose.yml not found"
    exit 1
fi

if ! docker compose config &> /dev/null; then
    print_fail "docker-compose.yml has syntax errors"
    exit 1
fi
print_pass "docker-compose.yml is valid"

# Test 2: Start services
print_test "Starting Docker Compose services"
if docker compose up -d &> /dev/null; then
    print_pass "Services started successfully"
    sleep 5  # Give services time to start
else
    print_fail "Failed to start services"
    exit 1
fi

# Test 3: Verify containers are running
print_test "Verifying containers are running"
RUNNING_CONTAINERS=$(docker compose ps --services --filter "status=running" | wc -l)
if [[ $RUNNING_CONTAINERS -ge 1 ]]; then
    print_pass "Services running: $RUNNING_CONTAINERS containers"
    docker compose ps
else
    print_fail "Not enough containers running"
    exit 1
fi

# Test 4: Check resource limits are applied
print_test "Verifying resource limits are applied to containers"
LIMITS_APPLIED=0

for service in $(docker compose config --services); do
    container_name="lab03-$service"
    if docker ps --format '{{.Names}}' | grep -q "$container_name"; then
        # Get CPU limit (NanoCPUs: 1e9 = 1 CPU core)
        nano_cpus=$(docker inspect "$container_name" --format '{{.HostConfig.NanoCpus}}' 2>/dev/null || echo "0")
        mem_limit=$(docker inspect "$container_name" --format '{{.HostConfig.Memory}}' 2>/dev/null || echo "0")

        # Convert NanoCPUs to CPU cores
        if [[ $nano_cpus -gt 0 ]]; then
            cpu_cores=$(echo "scale=2; $nano_cpus / 1000000000" | bc 2>/dev/null || echo "N/A")
        else
            cpu_cores="N/A"
        fi

        # Convert memory to GB/MB
        if [[ $mem_limit -gt 0 ]]; then
            if [[ $mem_limit -ge 1073741824 ]]; then
                mem_gb=$(echo "scale=2; $mem_limit / 1073741824" | bc 2>/dev/null || echo "N/A")
                mem_display="${mem_gb}G"
            else
                mem_mb=$(echo "scale=0; $mem_limit / 1048576" | bc 2>/dev/null || echo "N/A")
                mem_display="${mem_mb}M"
            fi
        else
            mem_display="N/A"
        fi

        if [[ "$cpu_cores" != "N/A" && "$mem_display" != "N/A" ]]; then
            print_pass "  $container_name: cpus: $cpu_cores, memory: $mem_display"
            ((LIMITS_APPLIED++)) || true
        else
            print_fail "  $container_name: Could not read limits (cpus: $cpu_cores, memory: $mem_display)"
        fi
    fi
done

if [[ $LIMITS_APPLIED -gt 0 ]]; then
    print_pass "Resource limits applied to $LIMITS_APPLIED containers"
else
    print_fail "No containers with resource limits found"
fi

# Test 5: Monitor resource usage with docker stats
print_test "Monitoring resource usage with docker stats"
STATS_OK=0

# Get initial stats
print_info "Current resource usage:"
docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.MemPerc}}" 2>/dev/null || true

for service in $(docker compose config --services); do
    container_name="lab03-$service"
    if docker ps --format '{{.Names}}' | grep -q "$container_name"; then
        # Get CPU percentage and memory usage
        cpu_percent=$(docker stats "$container_name" --no-stream --format "{{.CPUPerc}}" 2>/dev/null | sed 's/%//' || echo "N/A")
        mem_usage=$(docker stats "$container_name" --no-stream --format "{{.MemUsage}}" 2>/dev/null || echo "N/A")

        if [[ "$cpu_percent" != "N/A" && "$mem_usage" != "N/A" ]]; then
            print_info "  $container_name: CPU: ${cpu_percent}%, Memory: $mem_usage"
            ((STATS_OK++)) || true
        fi
    fi
done

if [[ $STATS_OK -gt 0 ]]; then
    print_pass "docker stats successfully monitored $STATS_OK containers"
else
    print_fail "Failed to monitor container stats"
fi

# Test 6: Verify health status
print_test "Verifying container health status"
HEALTHY_COUNT=0

for service in $(docker compose config --services); do
    container_name="lab03-$service"
    if docker ps --format '{{.Names}}' | grep -q "$container_name"; then
        # Get health status
        health_status=$(docker inspect "$container_name" --format '{{.State.Health.Status}}' 2>/dev/null || echo "no-healthcheck")

        if [[ "$health_status" == "healthy" ]]; then
            print_pass "  $container_name: $health_status"
            ((HEALTHY_COUNT++)) || true
        elif [[ "$health_status" == "no-healthcheck" ]]; then
            print_info "  $container_name: no healthcheck configured"
        else
            print_info "  $container_name: $health_status"
        fi
    fi
done

if [[ $HEALTHY_COUNT -gt 0 ]]; then
    print_pass "$HEALTHY_COUNT containers are healthy"
else
    print_warning "No healthy containers (services may still be starting)"
fi

# Test 7: Stress test for resource enforcement (if stress container exists)
print_test "Stress testing resource enforcement (if stress container available)"

STRESS_SERVICE=$(docker compose config --services | grep -E "(stress|test)" || echo "")
if [[ -n "$STRESS_SERVICE" ]]; then
    print_info "Found stress service: $STRESS_SERVICE"

    # Get the container name
    stress_container="lab03-$STRESS_SERVICE"
    if docker ps --format '{{.Names}}' | grep -q "$stress_container"; then
        # Get limits
        nano_cpus=$(docker inspect "$stress_container" --format '{{.HostConfig.NanoCpus}}' 2>/dev/null || echo "0")
        mem_limit=$(docker inspect "$stress_container" --format '{{.HostConfig.Memory}}' 2>/dev/null || echo "0")

        cpu_limit_cores=$(echo "scale=2; $nano_cpus / 1000000000" | bc 2>/dev/null || echo "N/A")
        if [[ $mem_limit -ge 1073741824 ]]; then
            mem_limit_display=$(echo "scale=2; $mem_limit / 1073741824" | bc 2>/dev/null || echo "N/A")"G"
        else
            mem_limit_display=$(echo "scale=0; $mem_limit / 1048576" | bc 2>/dev/null || echo "N/A")"M"
        fi

        print_info "Limits: cpus: $cpu_limit_cores, memory: $mem_limit_display"

        # Monitor during stress test
        print_info "Monitoring during stress (10 seconds)..."
        for i in {1..5}; do
            cpu_percent=$(docker stats "$stress_container" --no-stream --format "{{.CPUPerc}}" 2>/dev/null | sed 's/%//' || echo "N/A")
            mem_usage=$(docker stats "$stress_container" --no-stream --format "{{.MemUsage}}" 2>/dev/null || echo "N/A")
            print_info "  Sample $i: CPU: ${cpu_percent}%, Memory: $mem_usage"
            sleep 2
        done

        print_pass "Stress test completed"
    fi
else
    print_info "No stress service found - skipping stress test"
fi

# Test 8: Verify resource limit enforcement
print_test "Verifying resource limits are enforced"
ENFORCED=0

for service in $(docker compose config --services); do
    container_name="lab03-$service"
    if docker ps --format '{{.Names}}' | grep -q "$container_name"; then
        # Check if limits are set (not 0 or -1)
        nano_cpus=$(docker inspect "$container_name" --format '{{.HostConfig.NanoCpus}}' 2>/dev/null || echo "0")
        mem_limit=$(docker inspect "$container_name" --format '{{.HostConfig.Memory}}' 2>/dev/null || echo "0")

        if [[ $nano_cpus -gt 0 && $mem_limit -gt 0 ]]; then
            print_pass "  $container_name: Limits enforced (CPU and memory)"
            ((ENFORCED++)) || true
        else
            print_fail "  $container_name: Limits not enforced (CPU: $nano_cpus, Memory: $mem_limit)"
        fi
    fi
done

if [[ $ENFORCED -gt 0 ]]; then
    print_pass "Resource limits enforced for $ENFORCED containers"
else
    print_fail "Resource limits not enforced for any container"
fi

# Summary
print_header "Resource Enforcement Verification Summary"

echo -e "Tests run: $((pass_count + fail_count))"
echo -e "${GREEN}Passed: $pass_count${NC}"
if [[ $fail_count -gt 0 ]]; then
    echo -e "${RED}Failed: $fail_count${NC}"
fi

if [[ $fail_count -eq 0 ]]; then
    echo -e "\n${GREEN}${BOLD}✓ ALL ENFORCEMENT CHECKS PASSED${NC}"
    echo -e "\nResource limits are properly enforced!"
    exit 0
else
    echo -e "\n${RED}Some enforcement checks failed${NC}"
    echo -e "Please verify resource limit configuration."
    exit 1
fi