#!/bin/bash

# Memory profiling script for go-jdenticon
# This script provides easy-to-use commands for profiling memory usage

set -e

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
PROFILES_DIR="$PROJECT_ROOT/profiles"

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

# Ensure profiles directory exists
mkdir -p "$PROFILES_DIR"

usage() {
    cat << EOF
Usage: $0 [COMMAND] [OPTIONS]

Memory profiling commands for go-jdenticon:

COMMANDS:
    benchmark       Run benchmarks with memory profiling
    load-test       Run sustained load test with profiling
    analyze-heap    Analyze existing heap profile
    analyze-cpu     Analyze existing CPU profile
    web-ui          Start web UI for profile analysis
    continuous      Run continuous profiling session
    compare         Compare two profile files

OPTIONS:
    --duration=5m   Test duration (default: 5m)
    --workers=4     Number of workers (default: 4) 
    --cache=1000    Cache size (default: 1000)
    --output=FILE   Output profile file
    --baseline=FILE Baseline profile for comparison

EXAMPLES:
    $0 benchmark --duration=10m
    $0 load-test --workers=8 --cache=500
    $0 analyze-heap profiles/heap-latest.prof
    $0 continuous --duration=1h
    $0 compare profiles/baseline.prof profiles/after-fix.prof

PROFILING WORKFLOW:
    1. Run baseline: $0 benchmark --output=baseline
    2. Make changes to code
    3. Run comparison: $0 benchmark --output=after-changes
    4. Compare results: $0 compare profiles/baseline.prof profiles/after-changes.prof

EOF
}

log() {
    echo -e "${GREEN}[$(date +'%H:%M:%S')]${NC} $1"
}

warn() {
    echo -e "${YELLOW}[$(date +'%H:%M:%S')] WARNING:${NC} $1"
}

error() {
    echo -e "${RED}[$(date +'%H:%M:%S')] ERROR:${NC} $1"
    exit 1
}

# Parse command line arguments
COMMAND=""
DURATION="5m"
WORKERS="4"
CACHE_SIZE="1000"
OUTPUT_FILE=""
BASELINE_FILE=""

while [[ $# -gt 0 ]]; do
    case $1 in
        benchmark|load-test|analyze-heap|analyze-cpu|web-ui|continuous|compare)
            COMMAND="$1"
            shift
            ;;
        --duration=*)
            DURATION="${1#*=}"
            shift
            ;;
        --workers=*)
            WORKERS="${1#*=}"
            shift
            ;;
        --cache=*)
            CACHE_SIZE="${1#*=}"
            shift
            ;;
        --output=*)
            OUTPUT_FILE="${1#*=}"
            shift
            ;;
        --baseline=*)
            BASELINE_FILE="${1#*=}"
            shift
            ;;
        -h|--help)
            usage
            exit 0
            ;;
        *)
            if [[ -z "$COMMAND" ]]; then
                COMMAND="$1"
            else
                error "Unknown argument: $1"
            fi
            shift
            ;;
    esac
done

if [[ -z "$COMMAND" ]]; then
    usage
    exit 1
fi

# Generate timestamp for output files
TIMESTAMP=$(date +"%Y%m%d_%H%M%S")

run_benchmark() {
    log "Running memory leak validation benchmarks..."
    
    if [[ -n "$OUTPUT_FILE" ]]; then
        HEAP_PROF="$PROFILES_DIR/${OUTPUT_FILE}_heap.prof"
        CPU_PROF="$PROFILES_DIR/${OUTPUT_FILE}_cpu.prof"
        MEM_PROF="$PROFILES_DIR/${OUTPUT_FILE}_mem.prof"
    else
        HEAP_PROF="$PROFILES_DIR/benchmark_${TIMESTAMP}_heap.prof"
        CPU_PROF="$PROFILES_DIR/benchmark_${TIMESTAMP}_cpu.prof"
        MEM_PROF="$PROFILES_DIR/benchmark_${TIMESTAMP}_mem.prof"
    fi
    
    cd "$PROJECT_ROOT"
    
    log "Memory profiling benchmark with duration: $DURATION"
    log "Output files: heap=$HEAP_PROF, cpu=$CPU_PROF, mem=$MEM_PROF"
    
    # Run memory leak benchmarks with profiling
    go test -bench=BenchmarkMemoryLeak \
        -benchtime="$DURATION" \
        -memprofile="$MEM_PROF" \
        -memprofilerate=1 \
        -cpuprofile="$CPU_PROF" \
        -benchmem \
        ./jdenticon
    
    # Capture heap profile at the end
    if pgrep -f "go.*test.*BenchmarkMemoryLeak" > /dev/null; then
        warn "Test process still running, heap profile may not be available"
    fi
    
    log "Benchmark completed. Profile files saved:"
    log "  Memory: $MEM_PROF"
    log "  CPU: $CPU_PROF"
    
    # Create symlinks to latest profiles
    ln -sf "$(basename "$MEM_PROF")" "$PROFILES_DIR/mem-latest.prof"
    ln -sf "$(basename "$CPU_PROF")" "$PROFILES_DIR/cpu-latest.prof"
    
    log "Analysis commands:"
    log "  go tool pprof $MEM_PROF"
    log "  go tool pprof -http=:8080 $MEM_PROF"
}

run_load_test() {
    log "Running sustained load test with profiling..."
    
    cd "$PROJECT_ROOT"
    
    # Build the load test binary if it doesn't exist
    if [[ ! -f "cmd/memory-load-test/memory-load-test" ]]; then
        log "Building memory load test binary..."
        go build -o cmd/memory-load-test/memory-load-test ./cmd/memory-load-test
    fi
    
    HEAP_PROF="$PROFILES_DIR/loadtest_${TIMESTAMP}_heap.prof"
    CPU_PROF="$PROFILES_DIR/loadtest_${TIMESTAMP}_cpu.prof"
    
    log "Starting load test: duration=$DURATION, workers=$WORKERS, cache=$CACHE_SIZE"
    log "pprof server will be available at http://localhost:6060/debug/pprof/"
    log "Use Ctrl+C to stop the test gracefully"
    
    # Start the load test in background
    ./cmd/memory-load-test/memory-load-test \
        -duration="$DURATION" \
        -workers="$WORKERS" \
        -cache-size="$CACHE_SIZE" \
        -pprof-port=6060 &
    
    LOAD_TEST_PID=$!
    
    # Give it time to start
    sleep 5
    
    log "Load test started with PID $LOAD_TEST_PID"
    log "Monitoring will capture profiles every 30 seconds"
    
    # Wait for the test to complete or user interrupt
    trap "kill $LOAD_TEST_PID 2>/dev/null; wait $LOAD_TEST_PID 2>/dev/null" EXIT
    
    wait $LOAD_TEST_PID || true
    
    log "Load test completed"
}

analyze_heap() {
    local profile_file="$1"
    
    if [[ -z "$profile_file" ]]; then
        profile_file="$PROFILES_DIR/mem-latest.prof"
    fi
    
    if [[ ! -f "$profile_file" ]]; then
        error "Profile file not found: $profile_file"
    fi
    
    log "Analyzing heap profile: $profile_file"
    
    # Interactive analysis
    echo -e "\n${BLUE}Starting interactive pprof session...${NC}"
    echo -e "${BLUE}Useful commands:${NC}"
    echo "  top          - Show top memory allocations"
    echo "  top -cum     - Show top cumulative allocations"
    echo "  list <func>  - Show source code for function"
    echo "  web          - Open web interface"
    echo "  svg          - Generate SVG call graph"
    echo "  png          - Generate PNG call graph"
    echo "  traces       - Show allocation traces"
    echo "  help         - Show all commands"
    echo ""
    
    go tool pprof "$profile_file"
}

analyze_cpu() {
    local profile_file="$1"
    
    if [[ -z "$profile_file" ]]; then
        profile_file="$PROFILES_DIR/cpu-latest.prof"
    fi
    
    if [[ ! -f "$profile_file" ]]; then
        error "Profile file not found: $profile_file"
    fi
    
    log "Analyzing CPU profile: $profile_file"
    go tool pprof "$profile_file"
}

start_web_ui() {
    local profile_file="$1"
    
    if [[ -z "$profile_file" ]]; then
        profile_file="$PROFILES_DIR/mem-latest.prof"
    fi
    
    if [[ ! -f "$profile_file" ]]; then
        error "Profile file not found: $profile_file"
    fi
    
    log "Starting web UI for profile: $profile_file"
    log "Open http://localhost:8080 in your browser"
    
    go tool pprof -http=:8080 "$profile_file"
}

run_continuous() {
    log "Running continuous profiling session for $DURATION"
    
    cd "$PROJECT_ROOT"
    
    # Build load test if needed
    if [[ ! -f "cmd/memory-load-test/memory-load-test" ]]; then
        log "Building memory load test binary..."
        go build -o cmd/memory-load-test/memory-load-test ./cmd/memory-load-test
    fi
    
    log "Starting continuous load test..."
    
    # Start load test in background
    ./cmd/memory-load-test/memory-load-test \
        -duration="$DURATION" \
        -workers="$WORKERS" \
        -cache-size="$CACHE_SIZE" \
        -pprof-port=6060 \
        -varied-inputs=true &
    
    LOAD_TEST_PID=$!
    sleep 3
    
    log "Load test PID: $LOAD_TEST_PID"
    log "pprof endpoints available at http://localhost:6060/debug/pprof/"
    
    # Capture profiles at intervals
    INTERVAL=300  # 5 minutes
    PROFILES_CAPTURED=0
    
    while kill -0 $LOAD_TEST_PID 2>/dev/null; do
        sleep $INTERVAL
        
        PROFILES_CAPTURED=$((PROFILES_CAPTURED + 1))
        SNAPSHOT_TIME=$(date +"%H%M%S")
        
        log "Capturing profile snapshot $PROFILES_CAPTURED at $SNAPSHOT_TIME"
        
        # Capture heap profile
        curl -s "http://localhost:6060/debug/pprof/heap" > \
            "$PROFILES_DIR/continuous_${TIMESTAMP}_${PROFILES_CAPTURED}_heap.prof" || true
        
        # Capture goroutine profile
        curl -s "http://localhost:6060/debug/pprof/goroutine" > \
            "$PROFILES_DIR/continuous_${TIMESTAMP}_${PROFILES_CAPTURED}_goroutine.prof" || true
        
        log "  Saved heap and goroutine profiles"
    done
    
    log "Continuous profiling completed. Captured $PROFILES_CAPTURED snapshots."
    log "Profile files are in: $PROFILES_DIR/"
}

compare_profiles() {
    local baseline="$1"
    local current="$2"
    
    if [[ -z "$baseline" || -z "$current" ]]; then
        error "Usage: $0 compare <baseline_profile> <current_profile>"
    fi
    
    if [[ ! -f "$baseline" ]]; then
        error "Baseline profile not found: $baseline"
    fi
    
    if [[ ! -f "$current" ]]; then
        error "Current profile not found: $current"
    fi
    
    log "Comparing profiles:"
    log "  Baseline: $baseline"
    log "  Current:  $current"
    
    # Use pprof to compare profiles
    echo -e "\n${BLUE}Memory allocation differences:${NC}"
    go tool pprof -base="$baseline" -top "$current"
    
    echo -e "\n${BLUE}Interactive comparison session:${NC}"
    echo "Useful commands:"
    echo "  top          - Show top differences"
    echo "  top -cum     - Show cumulative differences"
    echo "  web          - Visual diff in browser"
    
    go tool pprof -base="$baseline" "$current"
}

# Execute the requested command
case "$COMMAND" in
    benchmark)
        run_benchmark
        ;;
    load-test)
        run_load_test
        ;;
    analyze-heap)
        analyze_heap "$2"
        ;;
    analyze-cpu)
        analyze_cpu "$2"
        ;;
    web-ui)
        start_web_ui "$2"
        ;;
    continuous)
        run_continuous
        ;;
    compare)
        compare_profiles "$2" "$3"
        ;;
    *)
        error "Unknown command: $COMMAND"
        ;;
esac