#!/usr/bin/env bash
set -e

echo "# Running benchmarks..."
echo ""

BENCH_DIR="$(pwd)/tests/benchmark"
IRC_BIN="$(pwd)/build/nix-irc"

GREEN='\033[0;32m'
BLUE='\033[0;34m'
YELLOW='\033[0;33m'
NC='\033[0m'

get_ms() {
	local time_str="$1"
	if [[ $time_str =~ ([0-9]+)m([0-9.]+)s ]]; then
		local mins="${BASH_REMATCH[1]}"
		local secs="${BASH_REMATCH[2]}"
		local ms
		ms=$(awk "BEGIN {printf \"%.1f\", ($mins * 60000) + ($secs * 1000)}")
		echo "$ms"
	else
		echo "0"
	fi
}

run_benchmark() {
	local name="$1"
	local file="$2"

	echo -e "${BLUE}=== $name ===${NC}"
	echo ""

	# Measure compilation time
	echo -n "  Compilation time:     "
	local compile_start
	compile_start=$(date +%s%N)
	"$IRC_BIN" "$file" /tmp/bench.nixir >/dev/null 2>&1
	local compile_end
	compile_end=$(date +%s%N)
	local compile_ms=$(((compile_end - compile_start) / 1000000))
	echo -e "${YELLOW}${compile_ms}ms${NC}"

	# Source and IR sizes
	local src_size
	src_size=$(stat -c%s "$file" 2>/dev/null || stat -f%z "$file" 2>/dev/null)
	local ir_size
	ir_size=$(stat -c%s /tmp/bench.nixir 2>/dev/null || stat -f%z /tmp/bench.nixir 2>/dev/null)
	local ratio=$((ir_size * 100 / src_size))
	echo -e "  Source size:          ${src_size}B"
	echo -e "  IR bundle size:       ${ir_size}B (${ratio}% of source)"

	echo ""

	# Native Nix evaluation (baseline)
	echo -n "  Native Nix eval:      "
	local native_total=0
	for _ in {1..5}; do
		local t
		t=$( (time nix-instantiate --eval --strict "$file" >/dev/null 2>&1) 2>&1 | grep "real" | awk '{print $2}')
		local ms
		ms=$(get_ms "$t")
		native_total=$(awk "BEGIN {print $native_total + $ms}")
	done
	local native_avg
	native_avg=$(awk "BEGIN {printf \"%.1f\", $native_total / 5}")
	echo -e "${GREEN}${native_avg}ms${NC} avg (5 runs)"

	echo ""
}

echo "Measuring IR compilation speed and bundle size characteristics."
echo ""

run_benchmark "Simple Expression" "$BENCH_DIR/simple.nix"
run_benchmark "Medium Complexity" "$BENCH_DIR/medium.nix"
run_benchmark "Large Expression" "$BENCH_DIR/large.nix"

# Overall statistics
echo -e "${BLUE}=== Overall Statistics ===${NC}"
echo ""

testdir=$(mktemp -d)
total_nix=0
total_ir=0
total_compile_time=0

for f in "$BENCH_DIR"/*.nix; do
	nixsize=$(stat -c%s "$f" 2>/dev/null || stat -f%z "$f" 2>/dev/null)
	base=$(basename "$f" .nix)
	irfile="${testdir}/${base}.nixir"

	start=$(date +%s%N)
	"$IRC_BIN" "$f" "$irfile" >/dev/null 2>&1
	end=$(date +%s%N)
	compile_time=$(((end - start) / 1000000))

	if [ -f "$irfile" ]; then
		irsize=$(stat -c%s "$irfile" 2>/dev/null || stat -f%z "$irfile" 2>/dev/null)
		total_nix=$((total_nix + nixsize))
		total_ir=$((total_ir + irsize))
		total_compile_time=$((total_compile_time + compile_time))
	fi
done

total_ratio=$((total_ir * 100 / total_nix))
avg_compile_time=$((total_compile_time / 3))

# TBH those are entirely unnecessary. However, I'm a sucker for data
# and those are trivial to compile. Might as well. Who knows, maybe it'll
# come in handy in the future.
echo "  Total source size:     ${total_nix}B"
echo "  Total IR size:         ${total_ir}B"
echo "  Compression ratio:     ${total_ratio}% of source"
echo "  Average compile time:  ${avg_compile_time}ms"
echo ""

rm -rf "$testdir"