#!/bin/bash

# Comprehensive performance benchmark script

TEST_FILE="${1:-documents/test.pdf}"

if [ ! -f "$TEST_FILE" ]; then
    echo "Error: Test file not found: $TEST_FILE"
    echo "Usage: $0 [pdf_file]"
    exit 1
fi

echo "========================================"
echo "Docling PDF Converter Performance Test"
echo "Test file: $TEST_FILE"
echo "========================================"

# Function to run test with monitoring
run_test() {
    local workers=$1
    local test_name="workers_${workers}"
    local output_dir="output_benchmark_${test_name}"

    echo ""
    echo "Testing with $workers workers"
    echo "--------------------------------"

    # Clean output directory
    rm -rf "$output_dir"

    # Start GPU monitoring
    "$(dirname "$0")/monitor_gpu.sh" "gpu_metrics_${test_name}.log" &
    MONITOR_PID=$!

    # Run conversion and time it
    START_TIME=$(date +%s)
    python src/converter.py "$TEST_FILE" --output-dir "$output_dir" --workers "$workers"
    EXIT_CODE=$?
    END_TIME=$(date +%s)
    DURATION=$((END_TIME - START_TIME))

    # Stop monitoring
    kill $MONITOR_PID 2>/dev/null
    wait $MONITOR_PID 2>/dev/null

    # Analyze results
    echo "Duration: ${DURATION} seconds"

    if [ -f "gpu_metrics_${test_name}.log" ]; then
        echo -n "Avg GPU: "
        awk -F', ' 'NR>1 {sum+=$4; count++} END {if (count>0) printf "%.1f%%\n", sum/count}' "gpu_metrics_${test_name}.log"
        echo -n "Peak GPU: "
        awk -F', ' 'NR>1 {if ($4>max) max=$4} END {printf "%.0f%%\n", max}' "gpu_metrics_${test_name}.log"
        echo -n "Avg Memory: "
        awk -F', ' 'NR>1 {sum+=$5; count++} END {if (count>0) printf "%.1f%%\n", sum/count}' "gpu_metrics_${test_name}.log"
        echo -n "Peak VRAM: "
        awk -F', ' 'NR>1 {if ($6>max) max=$6} END {printf "%d MB\n", max}' "gpu_metrics_${test_name}.log"
    fi

    # Count output files
    if [ -d "$output_dir" ]; then
        FILE_COUNT=$(ls -1 "$output_dir"/*.md 2>/dev/null | wc -l)
        echo "Output files: $FILE_COUNT"
    fi

    return $EXIT_CODE
}

# Test different worker configurations
for workers in 1 2 4; do
    run_test $workers
    sleep 5  # Allow GPU to cool down between tests
done

echo ""
echo "========================================"
echo "Benchmark Complete!"
echo ""
echo "Summary:"
echo "--------"

# Create summary table
printf "%-10s %-12s %-12s %-12s %-12s\n" "Workers" "Duration" "Avg GPU" "Peak GPU" "Peak VRAM"
printf "%-10s %-12s %-12s %-12s %-12s\n" "-------" "--------" "-------" "--------" "---------"

for workers in 1 2 4; do
    if [ -f "gpu_metrics_workers_${workers}.log" ]; then
        DURATION=$(grep -E "^Duration:" <<< "$(run_test $workers)" | cut -d' ' -f2)
        AVG_GPU=$(awk -F', ' 'NR>1 {sum+=$4; count++} END {if (count>0) printf "%.1f%%", sum/count}' "gpu_metrics_workers_${workers}.log")
        PEAK_GPU=$(awk -F', ' 'NR>1 {if ($4>max) max=$4} END {printf "%.0f%%", max}' "gpu_metrics_workers_${workers}.log")
        PEAK_VRAM=$(awk -F', ' 'NR>1 {if ($6>max) max=$6} END {printf "%d MB", max}' "gpu_metrics_workers_${workers}.log")
        printf "%-10s %-12s %-12s %-12s %-12s\n" "$workers" "${DURATION}s" "$AVG_GPU" "$PEAK_GPU" "$PEAK_VRAM"
    fi
done

echo ""
echo "Metrics saved to: gpu_metrics_workers_*.log"
echo "========================================"