#!/bin/bash

# Memory Leak Detection Script for go-jdenticon
# This script runs extended memory testing and captures heap profiles for leak analysis

set -e

# Configuration
TEST_DURATION="15m"
WORKERS=8
CACHE_SIZE=1000
PPROF_PORT=6060
PROFILE_DIR="./profiles"
SNAPSHOT_INTERVAL="2m"

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

echo -e "${BLUE}=== go-jdenticon Memory Leak Detection Test ===${NC}"
echo "Test duration: $TEST_DURATION"
echo "Workers: $WORKERS"
echo "Cache size: $CACHE_SIZE"
echo "Profile interval: $SNAPSHOT_INTERVAL"
echo ""

# Create profiles directory
mkdir -p "$PROFILE_DIR"
cd "$PROFILE_DIR" && rm -f *.prof *.json && cd ..

# Build the memory load test tool
echo -e "${YELLOW}Building memory-load-test tool...${NC}"
go build -o memory-load-test cmd/memory-load-test/main.go

# Start the load test in background
echo -e "${YELLOW}Starting extended memory load test...${NC}"
./memory-load-test \
  -duration="$TEST_DURATION" \
  -workers="$WORKERS" \
  -cache-size="$CACHE_SIZE" \
  -pprof-port="$PPROF_PORT" \
  -sample-interval="30s" \
  -report-interval="1m" \
  -varied-inputs=true &

LOAD_TEST_PID=$!

# Function to cleanup on exit
cleanup() {
  echo -e "\n${YELLOW}Cleaning up...${NC}"
  if kill -0 $LOAD_TEST_PID 2>/dev/null; then
    kill $LOAD_TEST_PID
    wait $LOAD_TEST_PID 2>/dev/null || true
  fi
  rm -f memory-load-test
}
trap cleanup EXIT

# Wait for pprof server to start
echo -e "${YELLOW}Waiting for pprof server to start...${NC}"
sleep 5

# Function to capture heap profile
capture_heap_profile() {
  local suffix=$1
  local gc_param=""
  
  # Force GC for some snapshots to get clean measurements
  if [[ "$suffix" == *"gc"* ]]; then
    gc_param="?gc=1"
  fi
  
  echo -e "${BLUE}Capturing heap profile: $suffix${NC}"
  curl -s "http://localhost:$PPROF_PORT/debug/pprof/heap$gc_param" > "$PROFILE_DIR/heap_$suffix.prof"
  
  # Also capture goroutine profile to check for goroutine leaks
  curl -s "http://localhost:$PPROF_PORT/debug/pprof/goroutine" > "$PROFILE_DIR/goroutine_$suffix.prof"
  
  # Get current memory stats from pprof
  curl -s "http://localhost:$PPROF_PORT/debug/pprof/heap?debug=1" | head -20 > "$PROFILE_DIR/memstats_$suffix.txt"
}

# Function to wait and show progress
wait_with_progress() {
  local duration=$1
  local message=$2
  local seconds=$(echo "$duration" | sed 's/m$//' | awk '{print $1*60}')
  
  echo -e "${YELLOW}$message (${duration})...${NC}"
  for ((i=1; i<=seconds; i++)); do
    if ((i % 30 == 0)); then
      echo -e "${BLUE}Progress: ${i}s / ${seconds}s${NC}"
    fi
    sleep 1
    
    # Check if load test is still running
    if ! kill -0 $LOAD_TEST_PID 2>/dev/null; then
      echo -e "${RED}Load test process died unexpectedly!${NC}"
      exit 1
    fi
  done
}

# Capture initial baseline profiles
echo -e "${GREEN}=== Capturing Baseline Profiles ===${NC}"
capture_heap_profile "0min_baseline"

# Wait and capture profiles at intervals
profiles_captured=1

# 2 minute mark
wait_with_progress "2m" "Waiting for 2-minute mark"
capture_heap_profile "2min"
((profiles_captured++))

# 5 minute mark  
wait_with_progress "3m" "Waiting for 5-minute mark"
capture_heap_profile "5min_gc"
((profiles_captured++))

# 8 minute mark
wait_with_progress "3m" "Waiting for 8-minute mark"
capture_heap_profile "8min"
((profiles_captured++))

# 10 minute mark (with GC)
wait_with_progress "2m" "Waiting for 10-minute mark"
capture_heap_profile "10min_gc"
((profiles_captured++))

# 13 minute mark
wait_with_progress "3m" "Waiting for 13-minute mark"
capture_heap_profile "13min"
((profiles_captured++))

# Wait for test completion and capture final profile
echo -e "${BLUE}Waiting for test completion...${NC}"
wait $LOAD_TEST_PID
LOAD_TEST_EXIT_CODE=$?

# Capture final profile after test completion
sleep 2
capture_heap_profile "15min_final"
((profiles_captured++))

echo -e "${GREEN}=== Memory Leak Analysis ===${NC}"
echo "Captured $profiles_captured heap profiles"

# Analyze profiles for memory leaks
echo -e "${YELLOW}Analyzing memory growth patterns...${NC}"

# Check if we have go tool pprof available
if ! command -v go &> /dev/null; then
  echo -e "${RED}Go toolchain not found - cannot analyze profiles${NC}"
  exit 1
fi

# Create leak analysis report
REPORT_FILE="$PROFILE_DIR/MEMORY_LEAK_ANALYSIS_$(date +%Y%m%d_%H%M%S).md"

cat > "$REPORT_FILE" << 'EOF'
# Memory Leak Analysis Report

## Test Configuration
- **Duration**: 15 minutes sustained load
- **Workers**: 8 concurrent goroutines
- **Cache Size**: 1000 entries
- **Input Variation**: Diverse input patterns to stress different code paths

## Profile Collection Points
1. **0min_baseline**: Initial state after startup
2. **2min**: Early runtime behavior
3. **5min_gc**: Mid-test with forced GC
4. **8min**: Sustained load behavior
5. **10min_gc**: Late-test with forced GC  
6. **13min**: Near end-of-test
7. **15min_final**: Final state after completion

EOF

# Function to analyze profile comparison
analyze_profile_diff() {
  local base_profile=$1
  local comp_profile=$2
  local period_name=$3
  
  echo -e "${BLUE}Analyzing: $base_profile vs $comp_profile ($period_name)${NC}"
  
  # Add section to report
  cat >> "$REPORT_FILE" << EOF

## $period_name Analysis

### Profile Comparison: $(basename "$base_profile" .prof) → $(basename "$comp_profile" .prof)

EOF

  # Generate diff analysis and save to report
  if go tool pprof -base="$base_profile" "$comp_profile" -top10 -cum >> "$REPORT_FILE" 2>/dev/null; then
    cat >> "$REPORT_FILE" << EOF

### Memory Growth Details
\`\`\`
EOF
    go tool pprof -base="$base_profile" "$comp_profile" -list=".*" 2>/dev/null | head -50 >> "$REPORT_FILE" || true
    cat >> "$REPORT_FILE" << EOF
\`\`\`

EOF
  else
    cat >> "$REPORT_FILE" << EOF

**⚠️ Profile analysis failed - may indicate inconsistent data or tool issues**

EOF
  fi
}

# Perform comparative analysis
analyze_profile_diff "$PROFILE_DIR/heap_0min_baseline.prof" "$PROFILE_DIR/heap_5min_gc.prof" "Early Growth (0-5 min)"
analyze_profile_diff "$PROFILE_DIR/heap_5min_gc.prof" "$PROFILE_DIR/heap_10min_gc.prof" "Mid Growth (5-10 min)"
analyze_profile_diff "$PROFILE_DIR/heap_10min_gc.prof" "$PROFILE_DIR/heap_15min_final.prof" "Late Growth (10-15 min)"
analyze_profile_diff "$PROFILE_DIR/heap_0min_baseline.prof" "$PROFILE_DIR/heap_15min_final.prof" "Overall Growth (0-15 min)"

# Analyze goroutine growth
echo -e "${YELLOW}Analyzing goroutine patterns...${NC}"

cat >> "$REPORT_FILE" << 'EOF'

## Goroutine Analysis

### Goroutine Count Over Time
EOF

for profile in "$PROFILE_DIR"/goroutine_*.prof; do
  if [[ -f "$profile" ]]; then
    timestamp=$(basename "$profile" .prof | sed 's/goroutine_//')
    count=$(go tool pprof -top "$profile" 2>/dev/null | grep "Total:" | awk '{print $2}' || echo "unknown")
    echo "- **$timestamp**: $count goroutines" >> "$REPORT_FILE"
  fi
done

# Memory stats analysis
echo -e "${YELLOW}Analyzing memory statistics...${NC}"

cat >> "$REPORT_FILE" << 'EOF'

## Memory Statistics Summary

### Heap Memory Usage Over Time
EOF

for memstat in "$PROFILE_DIR"/memstats_*.txt; do
  if [[ -f "$memstat" ]]; then
    timestamp=$(basename "$memstat" .txt | sed 's/memstats_//')
    heap_inuse=$(grep "# HeapInuse" "$memstat" | awk '{print $3}' || echo "unknown")
    heap_objects=$(grep "# HeapObjects" "$memstat" | awk '{print $3}' || echo "unknown")
    echo "- **$timestamp**: HeapInuse=$heap_inuse, HeapObjects=$heap_objects" >> "$REPORT_FILE"
  fi
done

# Load test results analysis
if [[ $LOAD_TEST_EXIT_CODE -eq 0 ]]; then
  cat >> "$REPORT_FILE" << 'EOF'

## Load Test Results
✅ **Load test completed successfully** - No crashes or critical errors detected.

EOF
else
  cat >> "$REPORT_FILE" << EOF

## Load Test Results
❌ **Load test failed with exit code $LOAD_TEST_EXIT_CODE** - May indicate memory issues or resource exhaustion.

EOF
fi

# Generate conclusions
cat >> "$REPORT_FILE" << 'EOF'

## Analysis Conclusions

### Memory Leak Assessment

EOF

# Try to detect obvious memory leaks from profile sizes
baseline_size=$(stat -c%s "$PROFILE_DIR/heap_0min_baseline.prof" 2>/dev/null || echo "0")
final_size=$(stat -c%s "$PROFILE_DIR/heap_15min_final.prof" 2>/dev/null || echo "0")

if [[ $baseline_size -gt 0 && $final_size -gt 0 ]]; then
  size_growth_ratio=$(echo "scale=2; $final_size / $baseline_size" | bc -l 2>/dev/null || echo "unknown")
  cat >> "$REPORT_FILE" << EOF
- **Profile Size Growth**: ${size_growth_ratio}x (${baseline_size} → ${final_size} bytes)
EOF
  
  if (( $(echo "$size_growth_ratio > 3.0" | bc -l 2>/dev/null || echo "0") )); then
    cat >> "$REPORT_FILE" << EOF
- ⚠️  **High profile growth** - Indicates potential memory accumulation
EOF
  elif (( $(echo "$size_growth_ratio > 1.5" | bc -l 2>/dev/null || echo "0") )); then
    cat >> "$REPORT_FILE" << EOF
- ⚠️  **Moderate profile growth** - Monitor for sustained increase trends
EOF
  else
    cat >> "$REPORT_FILE" << EOF
- ✅ **Reasonable profile growth** - Within expected bounds for sustained load
EOF
  fi
fi

cat >> "$REPORT_FILE" << 'EOF'

### Recommendations

1. **Profile Inspection**: Use `go tool pprof -http=:8080 profiles/heap_15min_final.prof` for interactive analysis
2. **Comparative Analysis**: Run `go tool pprof -base=profiles/heap_0min_baseline.prof profiles/heap_15min_final.prof` for detailed diff
3. **Goroutine Monitoring**: Check `go tool pprof profiles/goroutine_15min_final.prof` for goroutine leaks
4. **Long-term Testing**: Consider running tests for 1+ hours to detect slow leaks

### Commands for Manual Analysis

```bash
# Interactive heap analysis
go tool pprof -http=:8080 profiles/heap_15min_final.prof

# Memory growth analysis
go tool pprof -base=profiles/heap_0min_baseline.prof profiles/heap_15min_final.prof

# Goroutine analysis
go tool pprof profiles/goroutine_15min_final.prof

# View memory stats
cat profiles/memstats_15min_final.txt
```

EOF

echo -e "${GREEN}=== Memory Leak Test Complete ===${NC}"
echo -e "${BLUE}Results:${NC}"
echo "  - Profiles captured: $profiles_captured"
echo "  - Profile directory: $PROFILE_DIR"
echo "  - Analysis report: $REPORT_FILE"
echo "  - Load test exit code: $LOAD_TEST_EXIT_CODE"

if [[ $LOAD_TEST_EXIT_CODE -eq 0 ]]; then
  echo -e "${GREEN}  ✅ Load test completed successfully${NC}"
else
  echo -e "${RED}  ❌ Load test failed (exit code: $LOAD_TEST_EXIT_CODE)${NC}"
fi

echo ""
echo -e "${YELLOW}Next steps:${NC}"
echo "1. Review the analysis report: cat $REPORT_FILE"
echo "2. Interactive analysis: go tool pprof -http=:8080 $PROFILE_DIR/heap_15min_final.prof"
echo "3. Compare profiles: go tool pprof -base=$PROFILE_DIR/heap_0min_baseline.prof $PROFILE_DIR/heap_15min_final.prof"

# Final cleanup handled by trap