#!/usr/bin/env bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
PROJECT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"
EQUIV_DIR="$SCRIPT_DIR/equivalence"
DIAG_DIR="$SCRIPT_DIR/diagnostics"
LUMA="${TMPDIR:-/tmp}/gpu_shader_bin"

if [[ ! -f "$LUMA" ]] || [[ "$PROJECT_DIR/main.odin" -nt "$LUMA" ]] || \
   find "$PROJECT_DIR" -name '*.odin' -newer "$LUMA" 2>/dev/null | grep -q .; then
	echo "Building shader compiler..."
	(cd "$PROJECT_DIR/.." && odin build tools/shader -out:"$LUMA") || { echo "FAIL: build failed"; exit 1; }
fi

PASS=0
FAIL=0
ERRORS=""

compare_ir() {
	name="$1"
	old="$EQUIV_DIR/${name}_old.luma"
	new="$EQUIV_DIR/${name}_new.luma"
	old_ir="$(mktemp /tmp/luma_${name}_old_XXXXXX.ir)"
	new_ir="$(mktemp /tmp/luma_${name}_new_XXXXXX.ir)"
	old_norm="$(mktemp /tmp/luma_${name}_old_XXXXXX.norm)"
	new_norm="$(mktemp /tmp/luma_${name}_new_XXXXXX.norm)"

	"$LUMA" dump-ir "$old" > "$old_ir"
	"$LUMA" dump-ir "$new" > "$new_ir"
	normalize_ir "$old_ir" "$old_norm"
	normalize_ir "$new_ir" "$new_norm"
	if diff_output=$(diff -u "$old_norm" "$new_norm" 2>&1); then
		PASS=$((PASS + 1))
	else
		FAIL=$((FAIL + 1))
		ERRORS="${ERRORS}FAIL: ${name} old/new IR mismatch\n${diff_output}\n\n"
	fi
	rm -f "$old_ir" "$new_ir" "$old_norm" "$new_norm"
}

normalize_ir() {
	input="$1"
	output="$2"
	python3 - "$input" "$output" <<'PY'
import sys

src, dst = sys.argv[1], sys.argv[2]
lines = open(src, encoding="utf-8").read().splitlines()
structs = []
other = []
i = 0
while i < len(lines):
	line = lines[i]
	if line.startswith("struct "):
		block = [line]
		i += 1
		while i < len(lines):
			block.append(lines[i])
			if lines[i] == "end":
				i += 1
				break
			i += 1
		structs.append("\n".join(block))
		continue
	other.append(line)
	i += 1
with open(dst, "w", encoding="utf-8") as f:
	for block in sorted(structs):
		f.write(block)
		f.write("\n\n")
	for line in other:
		f.write(line)
		f.write("\n")
PY
}

expect_diag() {
	file="$1"
	pattern="$2"
	output="$("$LUMA" check "$DIAG_DIR/$file" 2>&1 || true)"
	if printf '%s\n' "$output" | grep -Fq "$pattern"; then
		PASS=$((PASS + 1))
	else
		FAIL=$((FAIL + 1))
		ERRORS="${ERRORS}FAIL: ${file} missing diagnostic: ${pattern}\n${output}\n\n"
	fi
}

compare_ir fragment
compare_ir vertex
compare_ir group

expect_diag missing_branch_output.luma "out slot 'color' is not definitely assigned on every path"
expect_diag mixed_return_out.luma "cannot mix explicit out writes with fragment return sugar"
expect_diag nested_return.luma "workflow stage return sugar is only supported as a top-level"
expect_diag wrong_stage_builtin.luma "builtin 'work.global_id' is only valid in compute stages"

if [[ $FAIL -gt 0 ]]; then
	echo -e "$ERRORS"
fi
echo "Workflow equivalence results: $PASS passed, $FAIL failed"
[[ $FAIL -eq 0 ]] || exit 1