#!/bin/bash
# Memory Cap Verification Script
# Tests that memory cap enforcement is working correctly

set -e

# Configuration
TEST_DIR="/tmp/mcp-vector-search-memory-test"
MAX_MEMORY_GB=2  # Use small limit for testing
CHECK_INTERVAL=2  # Check every 2 seconds

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

echo "======================================"
echo "Memory Cap Verification Test"
echo "======================================"
echo ""
echo "Test configuration:"
echo "  Memory cap: ${MAX_MEMORY_GB}GB"
echo "  Check interval: ${CHECK_INTERVAL}s"
echo ""

# Create test directory
rm -rf "$TEST_DIR"
mkdir -p "$TEST_DIR"

# Generate large test files to stress memory
echo "Generating test files..."
for i in {1..100}; do
    cat > "$TEST_DIR/test_$i.py" <<EOF
# Large test file $i
def test_function_$i():
    """Test function with lots of content."""
    data = [
        "This is line 1 with lots of content" * 100,
        "This is line 2 with lots of content" * 100,
        "This is line 3 with lots of content" * 100,
    ] * 100
    return data

class TestClass$i:
    """Test class with lots of methods."""
    def __init__(self):
        self.data = "x" * 10000

    def method_1(self):
        return "method_1" * 100

    def method_2(self):
        return "method_2" * 100

    def method_3(self):
        return "method_3" * 100
EOF
done

echo "Generated 100 test files in $TEST_DIR"
echo ""

# Start indexing in background
echo "Starting indexing with ${MAX_MEMORY_GB}GB memory cap..."
export MCP_VECTOR_SEARCH_MAX_MEMORY_GB=$MAX_MEMORY_GB

# Start indexing in background
mcp-vector-search index "$TEST_DIR" > /tmp/mcp-memory-test.log 2>&1 &
INDEXER_PID=$!

echo "Indexer started (PID: $INDEXER_PID)"
echo ""

# Monitor memory usage
MAX_MEMORY_SEEN=0
MEMORY_VIOLATIONS=0
CHECKS=0

echo "Monitoring memory usage..."
echo "Press Ctrl+C to stop"
echo ""

while kill -0 $INDEXER_PID 2>/dev/null; do
    # Get memory usage in GB
    MEMORY_MB=$(ps -o rss= -p $INDEXER_PID 2>/dev/null || echo "0")
    MEMORY_GB=$(echo "scale=2; $MEMORY_MB / 1024 / 1024" | bc)

    # Track max memory
    if (( $(echo "$MEMORY_GB > $MAX_MEMORY_SEEN" | bc -l) )); then
        MAX_MEMORY_SEEN=$MEMORY_GB
    fi

    # Check for violations
    if (( $(echo "$MEMORY_GB > $MAX_MEMORY_GB" | bc -l) )); then
        MEMORY_VIOLATIONS=$((MEMORY_VIOLATIONS + 1))
        echo -e "${RED}⚠️  VIOLATION: Memory usage ${MEMORY_GB}GB exceeds cap of ${MAX_MEMORY_GB}GB${NC}"
    else
        echo -e "${GREEN}✓${NC} Memory: ${MEMORY_GB}GB / ${MAX_MEMORY_GB}GB (${MAX_MEMORY_SEEN}GB peak)"
    fi

    CHECKS=$((CHECKS + 1))
    sleep $CHECK_INTERVAL
done

echo ""
echo "======================================"
echo "Test Results"
echo "======================================"
echo "Total checks: $CHECKS"
echo "Peak memory: ${MAX_MEMORY_SEEN}GB"
echo "Memory cap: ${MAX_MEMORY_GB}GB"
echo "Violations: $MEMORY_VIOLATIONS"
echo ""

# Check logs for memory warnings
MEMORY_WARNINGS=$(grep -c "Memory limit exceeded" /tmp/mcp-memory-test.log || true)
BATCH_REDUCTIONS=$(grep -c "Reduced batch size" /tmp/mcp-memory-test.log || true)

echo "Log analysis:"
echo "  Memory warnings: $MEMORY_WARNINGS"
echo "  Batch reductions: $BATCH_REDUCTIONS"
echo ""

# Final verdict
if [ $MEMORY_VIOLATIONS -eq 0 ]; then
    echo -e "${GREEN}✅ PASS: Memory cap enforced correctly${NC}"
    echo "Peak memory (${MAX_MEMORY_SEEN}GB) stayed within limit (${MAX_MEMORY_GB}GB)"
    exit 0
else
    echo -e "${RED}❌ FAIL: Memory cap violated $MEMORY_VIOLATIONS times${NC}"
    echo "Peak memory (${MAX_MEMORY_SEEN}GB) exceeded limit (${MAX_MEMORY_GB}GB)"
    echo ""
    echo "Recent log entries:"
    tail -20 /tmp/mcp-memory-test.log
    exit 1
fi
