#!/bin/bash
# GDELT Engine Stress Test Suite
# Tests API endpoints with various load scenarios and monitors system resources

set -e

API_URL="${API_URL:-http://localhost:8080}"
MONITOR_INTERVAL=1

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

# Generate valid GDELT timestamps (must be 00, 15, 30, or 45 minutes)
generate_timestamps() {
    local count=$1
    local base_date="${2:-20260128}"
    local timestamps=()
    
    hours=(17 18 19 20 21 22 23 00 01 02 03 04 05 06 07 08 09 10 11 12 13 14 15 16)
    minutes=(00 15 30 45)
    
    for ((i=0; i<count; i++)); do
        hour_idx=$((i / 4 % 24))
        min_idx=$((i % 4))
        hour=$(printf "%02d" ${hours[$hour_idx]})
        min=${minutes[$min_idx]}
        timestamps+=("${base_date}${hour}${min}00")
    done
    
    echo "${timestamps[@]}"
}

# Monitor system resources
monitor_resources() {
    local pid=$1
    local duration=$2
    local output_file=$3
    
    echo "timestamp,cpu_percent,mem_mb,goroutines" > "$output_file"
    
    for ((i=0; i<duration; i++)); do
        if kill -0 $pid 2>/dev/null; then
            # Get CPU and memory for the process
            stats=$(ps -p $pid -o %cpu,rss --no-headers 2>/dev/null || echo "0 0")
            cpu=$(echo $stats | awk '{print $1}')
            mem_kb=$(echo $stats | awk '{print $2}')
            mem_mb=$((mem_kb / 1024))
            
            # Get goroutine count if pprof is available
            goroutines=$(curl -s "${API_URL}/debug/pprof/goroutine?debug=0" 2>/dev/null | wc -l || echo "N/A")
            
            echo "$(date +%s),$cpu,$mem_mb,$goroutines" >> "$output_file"
        fi
        sleep $MONITOR_INTERVAL
    done
}

print_header() {
    echo ""
    echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}"
    echo -e "${BLUE}  $1${NC}"
    echo -e "${BLUE}═══════════════════════════════════════════════════════════${NC}"
}

print_result() {
    local name=$1
    local status=$2
    local duration=$3
    
    if [ "$status" = "PASS" ]; then
        echo -e "${GREEN}✓ $name${NC} - ${duration}ms"
    else
        echo -e "${RED}✗ $name${NC} - $status"
    fi
}

# ═══════════════════════════════════════════════════════════
# TEST 1: Health Check
# ═══════════════════════════════════════════════════════════
test_health() {
    print_header "Test 1: Health Check"
    
    start=$(date +%s%3N)
    response=$(curl -s -w "\n%{http_code}" "${API_URL}/health")
    http_code=$(echo "$response" | tail -1)
    body=$(echo "$response" | head -n -1)
    end=$(date +%s%3N)
    duration=$((end - start))
    
    if [ "$http_code" = "200" ]; then
        print_result "Health endpoint" "PASS" "$duration"
        echo "$body" | jq .
    else
        print_result "Health endpoint" "FAIL: HTTP $http_code" "$duration"
    fi
}

# ═══════════════════════════════════════════════════════════
# TEST 2: Stats Check
# ═══════════════════════════════════════════════════════════
test_stats() {
    print_header "Test 2: Database Stats"
    
    start=$(date +%s%3N)
    response=$(curl -s -w "\n%{http_code}" "${API_URL}/stats")
    http_code=$(echo "$response" | tail -1)
    body=$(echo "$response" | head -n -1)
    end=$(date +%s%3N)
    duration=$((end - start))
    
    if [ "$http_code" = "200" ]; then
        print_result "Stats endpoint" "PASS" "$duration"
        echo "$body" | jq .
    else
        print_result "Stats endpoint" "FAIL: HTTP $http_code" "$duration"
    fi
}

# ═══════════════════════════════════════════════════════════
# TEST 3: Validation Tests
# ═══════════════════════════════════════════════════════════
test_validation() {
    print_header "Test 3: Timestamp Validation"
    
    echo -e "\n${YELLOW}Testing invalid timestamps:${NC}"
    
    # Invalid length
    echo -n "  Invalid length (12 chars): "
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d '{"timestamps": ["202601281715"]}')
    echo "$response" | jq -c '.rejected[0].reason // .error'
    
    # Invalid minute
    echo -n "  Invalid minute (17): "
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d '{"timestamps": ["20260128171700"]}')
    echo "$response" | jq -c '.rejected[0].reason // .error'
    
    # Future timestamp
    echo -n "  Future timestamp: "
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d '{"timestamps": ["20990128171500"]}')
    echo "$response" | jq -c '.rejected[0].reason // .error'
    
    # Valid timestamp
    echo -e "\n${YELLOW}Testing valid timestamp:${NC}"
    echo -n "  Valid format: "
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d '{"timestamps": ["20260128171500"]}')
    echo "$response" | jq -c '{accepted, queued: .timestamps_queued, rejected: .timestamps_rejected}'
}

# ═══════════════════════════════════════════════════════════
# TEST 4: Single Timestamp Processing
# ═══════════════════════════════════════════════════════════
test_single() {
    print_header "Test 4: Single Timestamp Processing"
    
    local ts="20260128180000"
    
    echo -e "${YELLOW}Submitting 1 timestamp: $ts${NC}"
    
    start=$(date +%s%3N)
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d "{\"timestamps\": [\"$ts\"]}")
    end=$(date +%s%3N)
    
    echo "$response" | jq .
    
    echo -e "\n${YELLOW}Waiting for completion...${NC}"
    
    for i in {1..60}; do
        sleep 2
        status=$(curl -s "${API_URL}/status/$ts")
        current_status=$(echo "$status" | jq -r '.status')
        
        if [ "$current_status" = "completed" ]; then
            echo -e "${GREEN}✓ Completed!${NC}"
            echo "$status" | jq .
            break
        elif [ "$current_status" = "failed" ] || [ "$current_status" = "invalid" ]; then
            echo -e "${RED}✗ Failed!${NC}"
            echo "$status" | jq .
            break
        else
            echo -n "."
        fi
    done
}

# ═══════════════════════════════════════════════════════════
# TEST 5: Batch Processing (4 timestamps)
# ═══════════════════════════════════════════════════════════
test_batch_small() {
    print_header "Test 5: Small Batch (4 timestamps)"
    
    timestamps=$(generate_timestamps 4 "20260127")
    ts_array=$(echo $timestamps | tr ' ' ',' | sed 's/^/["/;s/,/","/g;s/$/"]/')
    
    echo -e "${YELLOW}Submitting: $ts_array${NC}"
    
    start=$(date +%s)
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d "{\"timestamps\": $ts_array}")
    
    echo "$response" | jq -c '{accepted, queued: .timestamps_queued}'
    
    echo -e "\n${YELLOW}Monitoring progress...${NC}"
    
    for i in {1..120}; do
        sleep 2
        list=$(curl -s "${API_URL}/timestamps")
        completed=$(echo "$list" | jq '.total_completed')
        processing=$(echo "$list" | jq '.total_processing')
        
        echo -ne "\r  Completed: $completed | Processing: $processing    "
        
        if [ "$processing" = "0" ]; then
            break
        fi
    done
    
    end=$(date +%s)
    echo -e "\n${GREEN}✓ Batch completed in $((end - start)) seconds${NC}"
}

# ═══════════════════════════════════════════════════════════
# TEST 6: Large Batch Stress Test
# ═══════════════════════════════════════════════════════════
test_batch_large() {
    print_header "Test 6: Large Batch Stress Test (24 timestamps)"
    
    timestamps=$(generate_timestamps 24 "20260126")
    ts_array=$(echo $timestamps | tr ' ' ',' | sed 's/^/["/;s/,/","/g;s/$/"]/')
    
    echo -e "${YELLOW}Submitting 24 timestamps...${NC}"
    
    # Start resource monitoring in background
    if [ -n "$ENGINE_PID" ]; then
        monitor_resources $ENGINE_PID 300 "/tmp/gdelt_stress_resources.csv" &
        MONITOR_PID=$!
    fi
    
    start=$(date +%s)
    response=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d "{\"timestamps\": $ts_array}")
    
    echo "$response" | jq -c '{accepted, queued: .timestamps_queued}'
    
    echo -e "\n${YELLOW}Monitoring progress (this may take 2-3 minutes)...${NC}"
    
    for i in {1..180}; do
        sleep 2
        list=$(curl -s "${API_URL}/timestamps")
        completed=$(echo "$list" | jq '.total_completed')
        processing=$(echo "$list" | jq '.total_processing')
        
        # Get current system stats
        cpu=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | cut -d'%' -f1)
        mem=$(free -m | awk 'NR==2{printf "%.1f", $3/$2*100}')
        
        echo -ne "\r  Completed: $completed | Processing: $processing | CPU: ${cpu}% | RAM: ${mem}%    "
        
        if [ "$processing" = "0" ]; then
            break
        fi
    done
    
    end=$(date +%s)
    
    # Stop monitoring
    if [ -n "$MONITOR_PID" ]; then
        kill $MONITOR_PID 2>/dev/null || true
    fi
    
    echo -e "\n${GREEN}✓ Stress test completed in $((end - start)) seconds${NC}"
    
    # Show final stats
    echo -e "\n${YELLOW}Final Statistics:${NC}"
    curl -s "${API_URL}/stats" | jq .
}

# ═══════════════════════════════════════════════════════════
# TEST 7: Duplicate/Idempotency Test
# ═══════════════════════════════════════════════════════════
test_idempotency() {
    print_header "Test 7: Idempotency Test"
    
    local ts="20260128183000"
    
    echo -e "${YELLOW}Submitting same timestamp twice:${NC}"
    
    echo -n "  First request: "
    response1=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d "{\"timestamps\": [\"$ts\"]}")
    echo "$response1" | jq -c '{queued: .timestamps_queued, rejected: .timestamps_rejected}'
    
    sleep 1
    
    echo -n "  Second request (should reject): "
    response2=$(curl -s "${API_URL}/process" -X POST -H "Content-Type: application/json" \
        -d "{\"timestamps\": [\"$ts\"]}")
    echo "$response2" | jq -c '{queued: .timestamps_queued, rejected: .timestamps_rejected}'
}

# ═══════════════════════════════════════════════════════════
# MAIN
# ═══════════════════════════════════════════════════════════

echo ""
echo -e "${GREEN}╔═══════════════════════════════════════════════════════════╗${NC}"
echo -e "${GREEN}║         GDELT Engine Stress Test Suite                    ║${NC}"
echo -e "${GREEN}╚═══════════════════════════════════════════════════════════╝${NC}"
echo ""
echo -e "API URL: ${BLUE}${API_URL}${NC}"
echo ""

# Find engine PID if running
ENGINE_PID=$(pgrep -f "go run cmd/main.go" 2>/dev/null || pgrep -f "./engine" 2>/dev/null || echo "")
if [ -n "$ENGINE_PID" ]; then
    echo -e "Engine PID: ${BLUE}${ENGINE_PID}${NC}"
fi

# Parse arguments
case "${1:-all}" in
    health)
        test_health
        ;;
    stats)
        test_stats
        ;;
    validate)
        test_validation
        ;;
    single)
        test_single
        ;;
    batch)
        test_batch_small
        ;;
    stress)
        test_batch_large
        ;;
    idempotency)
        test_idempotency
        ;;
    all)
        test_health
        test_stats
        test_validation
        # Comment out heavy tests by default
        # test_single
        # test_batch_small
        # test_batch_large
        # test_idempotency
        echo -e "\n${YELLOW}Quick tests complete. Run individual tests for processing:${NC}"
        echo "  ./scripts/stress_test.sh single      # Test 1 timestamp"
        echo "  ./scripts/stress_test.sh batch       # Test 4 timestamps"
        echo "  ./scripts/stress_test.sh stress      # Test 24 timestamps"
        echo "  ./scripts/stress_test.sh idempotency # Test duplicate handling"
        ;;
    *)
        echo "Usage: $0 {health|stats|validate|single|batch|stress|idempotency|all}"
        exit 1
        ;;
esac

echo ""
echo -e "${GREEN}Done!${NC}"