#!/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 only
	echo -n "  Compilation only:     "
	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}"

	# Measure IR loading only (deserialization + evaluation)
	echo -n "  IR load only:         "
	PLUGIN_PATH="$(pwd)/build/nix-ir-plugin.so"
	if [ ! -f "$PLUGIN_PATH" ]; then
		echo -e "${YELLOW}skipped${NC} (plugin not built)"
	else
		# Pre-compile the IR
		"$IRC_BIN" "$file" /tmp/bench.nixir >/dev/null 2>&1

		# Measure just the loading (average of 10 runs to reduce noise)
		local total_load_us=0
		for _ in {1..10}; do
			local load_output
			load_output=$(nix-instantiate --plugin-files "$PLUGIN_PATH" --eval --expr "builtins.nixIR_loadIR \"/tmp/bench.nixir\"" 2>&1 >/dev/null | grep "nixIR timing" | grep -oP 'total=\K[0-9]+')
			total_load_us=$((total_load_us + load_output))
		done
		local avg_load_us=$((total_load_us / 10))
		local avg_load_ms_frac=$(awk "BEGIN {printf \"%.3f\", $avg_load_us / 1000}")
		echo -e "${GREEN}${avg_load_ms_frac}ms${NC} avg (10 runs)"
	fi

	# Measure full pipeline (compile + nix-instantiate overhead + IR load)
	echo -n "  Full pipeline:        "
	if [ ! -f "$PLUGIN_PATH" ]; then
		echo -e "${YELLOW}skipped${NC}"
	else
		local pipeline_start
		pipeline_start=$(date +%s%N)
		"$IRC_BIN" "$file" /tmp/bench.nixir >/dev/null 2>&1
		nix-instantiate --plugin-files "$PLUGIN_PATH" --eval --expr "builtins.nixIR_loadIR \"/tmp/bench.nixir\"" >/dev/null 2>&1
		local pipeline_end
		pipeline_end=$(date +%s%N)
		local pipeline_ms=$(((pipeline_end - pipeline_start) / 1000000))
		echo -e "${YELLOW}${pipeline_ms}ms${NC}"
	fi

	# 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=0
	if [[ "$src_size" -gt 0 ]]; then
		ratio=$((ir_size * 100 / src_size))
	fi
	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"