#!/bin/bash
#
# Flame Graph Generation Script for Performance Profiling
#
# This script generates flame graphs for the Foxhunt HFT system to visualize
# where CPU time is spent during benchmark execution. Useful for identifying
# performance bottlenecks and optimization opportunities.
#
# Usage:
#   ./scripts/generate_flame_graphs.sh [benchmark-name] [duration-seconds]
#
# Examples:
#   ./scripts/generate_flame_graphs.sh                    # All benchmarks, 30s each
#   ./scripts/generate_flame_graphs.sh ml_prediction 60   # Specific benchmark, 60s
#
# Requirements:
#   - cargo-flamegraph (install with: cargo install flamegraph)
#   - perf (Linux kernel profiler)
#   - Linux system (perf required)

set -euo pipefail

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

# Configuration
BENCHMARK_NAME="${1:-all}"
PROFILE_DURATION="${2:-30}"
OUTPUT_DIR="flame_graphs"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)

echo -e "${BLUE}=================================================${NC}"
echo -e "${BLUE}Flame Graph Generation for Performance Profiling${NC}"
echo -e "${BLUE}=================================================${NC}"
echo ""
echo -e "${GREEN}Benchmark:${NC} $BENCHMARK_NAME"
echo -e "${GREEN}Duration:${NC} ${PROFILE_DURATION}s"
echo -e "${GREEN}Output:${NC} $OUTPUT_DIR"
echo ""

# Check if running on Linux
if [[ "$OSTYPE" != "linux-gnu"* ]]; then
    echo -e "${RED}Error: Flame graphs require Linux (perf tool)${NC}"
    echo -e "${YELLOW}Alternatives:${NC}"
    echo "  - Run in Docker: docker run -v \$(pwd):/workspace rust:latest bash"
    echo "  - Use Instruments on macOS"
    echo "  - Use other profiling tools (valgrind, etc.)"
    exit 1
fi

# Check if cargo-flamegraph is installed
if ! command -v cargo-flamegraph &> /dev/null; then
    echo -e "${YELLOW}Installing cargo-flamegraph...${NC}"
    if ! cargo install flamegraph; then
        echo -e "${RED}Failed to install cargo-flamegraph${NC}"
        exit 1
    fi
    echo -e "${GREEN}✓ cargo-flamegraph installed${NC}"
fi

# Check if perf is available
if ! command -v perf &> /dev/null; then
    echo -e "${YELLOW}perf not found. Attempting to install...${NC}"
    if sudo apt-get install -y linux-tools-common linux-tools-generic linux-tools-$(uname -r) 2>/dev/null; then
        echo -e "${GREEN}✓ perf installed${NC}"
    else
        echo -e "${RED}Failed to install perf. Please install manually:${NC}"
        echo "  sudo apt-get install linux-tools-\$(uname -r)"
        exit 1
    fi
fi

# Set perf permissions (may require sudo)
echo -e "${YELLOW}Configuring perf permissions...${NC}"
if [ -f /proc/sys/kernel/perf_event_paranoid ]; then
    CURRENT_PARANOID=$(cat /proc/sys/kernel/perf_event_paranoid)
    if [ "$CURRENT_PARANOID" -gt 1 ]; then
        echo -e "${YELLOW}Note: perf_event_paranoid is $CURRENT_PARANOID (restrictive)${NC}"
        echo -e "${YELLOW}For better profiling, set to -1 (requires sudo):${NC}"
        echo "  sudo sysctl kernel.perf_event_paranoid=-1"
        echo ""
        read -p "Attempt to set now? (y/N) " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            sudo sysctl kernel.perf_event_paranoid=-1
        fi
    fi
fi

# Create output directory
mkdir -p "$OUTPUT_DIR"

# Function to generate flame graph for specific benchmark
generate_flamegraph() {
    local bench_name="$1"
    local output_file="$OUTPUT_DIR/flamegraph_${bench_name}_${TIMESTAMP}.svg"

    echo ""
    echo -e "${YELLOW}Generating flame graph: $bench_name${NC}"
    echo -e "${BLUE}Duration: ${PROFILE_DURATION}s${NC}"
    echo ""

    # Run cargo-flamegraph on the benchmark
    if cargo flamegraph \
        --bench performance_regression \
        --output "$output_file" \
        -- --bench "$bench_name" --profile-time "$PROFILE_DURATION"; then

        echo -e "${GREEN}✓ Flame graph generated: $output_file${NC}"

        # Get file size
        local file_size=$(du -h "$output_file" | cut -f1)
        echo -e "${BLUE}  Size: $file_size${NC}"

        return 0
    else
        echo -e "${RED}✗ Failed to generate flame graph for $bench_name${NC}"
        return 1
    fi
}

# Main execution
main() {
    if [ "$BENCHMARK_NAME" = "all" ]; then
        echo -e "${YELLOW}Generating flame graphs for all benchmarks...${NC}"
        echo -e "${BLUE}This will take approximately $((PROFILE_DURATION * 8)) seconds${NC}"
        echo ""

        # List of all benchmarks in performance_regression.rs
        BENCHMARKS=(
            "ml_prediction_latency"
            "hot_swap_latency"
            "database_writes"
            "backtest_performance"
            "order_processing"
            "risk_validation"
            "memory_allocation"
            "concurrent_access"
        )

        local success_count=0
        local total_count=${#BENCHMARKS[@]}

        for bench in "${BENCHMARKS[@]}"; do
            if generate_flamegraph "$bench"; then
                ((success_count++))
            fi
        done

        echo ""
        echo -e "${GREEN}=================================================${NC}"
        echo -e "${GREEN}Flame Graph Generation Complete${NC}"
        echo -e "${GREEN}=================================================${NC}"
        echo -e "${GREEN}Success: $success_count / $total_count${NC}"
        echo ""
    else
        generate_flamegraph "$BENCHMARK_NAME"
    fi

    # Generate summary HTML
    echo ""
    echo -e "${YELLOW}Generating summary page...${NC}"

    cat > "$OUTPUT_DIR/index.html" <<'EOFHTML'
<!DOCTYPE html>
<html>
<head>
    <title>Foxhunt Performance Flame Graphs</title>
    <style>
        body {
            font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
            max-width: 1200px;
            margin: 0 auto;
            padding: 20px;
            background: #f5f5f5;
        }
        h1 {
            color: #333;
            border-bottom: 3px solid #4CAF50;
            padding-bottom: 10px;
        }
        .benchmark {
            background: white;
            border-radius: 8px;
            padding: 20px;
            margin: 20px 0;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .benchmark h2 {
            color: #4CAF50;
            margin-top: 0;
        }
        .benchmark iframe {
            width: 100%;
            height: 600px;
            border: 1px solid #ddd;
            border-radius: 4px;
        }
        .info {
            background: #e3f2fd;
            border-left: 4px solid #2196F3;
            padding: 15px;
            margin: 20px 0;
            border-radius: 4px;
        }
        .legend {
            background: #fff3e0;
            border-left: 4px solid #ff9800;
            padding: 15px;
            margin: 20px 0;
            border-radius: 4px;
        }
    </style>
</head>
<body>
    <h1>Foxhunt HFT Performance Flame Graphs</h1>

    <div class="info">
        <strong>Generated:</strong> TIMESTAMP_PLACEHOLDER<br>
        <strong>Profile Duration:</strong> DURATION_PLACEHOLDER seconds per benchmark
    </div>

    <div class="legend">
        <h3>How to Read Flame Graphs</h3>
        <ul>
            <li><strong>X-axis:</strong> Alphabetical ordering (NOT time)</li>
            <li><strong>Y-axis:</strong> Stack depth (call hierarchy)</li>
            <li><strong>Width:</strong> CPU time consumed</li>
            <li><strong>Color:</strong> Random (for differentiation only)</li>
        </ul>
        <p><strong>Tip:</strong> Click on any frame to zoom in. Look for wide frames at the top for optimization opportunities.</p>
    </div>

EOFHTML

    # Add each flame graph to the HTML
    for svg_file in "$OUTPUT_DIR"/flamegraph_*.svg; do
        if [ -f "$svg_file" ]; then
            local bench_name=$(basename "$svg_file" | sed 's/flamegraph_//;s/_[0-9]*\.svg$//')
            local svg_filename=$(basename "$svg_file")

            cat >> "$OUTPUT_DIR/index.html" <<EOFBENCH
    <div class="benchmark">
        <h2>$bench_name</h2>
        <iframe src="$svg_filename" scrolling="no"></iframe>
    </div>

EOFBENCH
        fi
    done

    cat >> "$OUTPUT_DIR/index.html" <<'EOFFOOTER'
</body>
</html>
EOFFOOTER

    # Replace placeholders
    sed -i "s/TIMESTAMP_PLACEHOLDER/$(date)/" "$OUTPUT_DIR/index.html"
    sed -i "s/DURATION_PLACEHOLDER/$PROFILE_DURATION/" "$OUTPUT_DIR/index.html"

    echo -e "${GREEN}✓ Summary page generated: $OUTPUT_DIR/index.html${NC}"

    # Show results
    echo ""
    echo -e "${GREEN}=================================================${NC}"
    echo -e "${GREEN}Flame Graphs Generated Successfully${NC}"
    echo -e "${GREEN}=================================================${NC}"
    echo ""
    echo -e "${BLUE}View results:${NC}"
    echo -e "  ${YELLOW}open $OUTPUT_DIR/index.html${NC}"
    echo ""
    echo -e "${BLUE}Individual flame graphs:${NC}"
    for svg_file in "$OUTPUT_DIR"/flamegraph_*.svg; do
        if [ -f "$svg_file" ]; then
            echo -e "  - $(basename "$svg_file")"
        fi
    done
    echo ""
    echo -e "${BLUE}Analysis tips:${NC}"
    echo "  1. Look for wide horizontal bars (high CPU time)"
    echo "  2. Deep stacks indicate complex call chains"
    echo "  3. Optimize functions with widest frames"
    echo "  4. Compare flame graphs before/after optimization"
    echo ""
}

# Run main function
main "$@"