#!/bin/bash
# run_tb.sh - Run Vivado XSIM testbench for a module
#
# Usage: ./run_tb.sh <module> [K] [CASE]
#        ./run_tb.sh ntt                  # run the module's full xsim_run.tcl
#        ./run_tb.sh top                  # ML-KEM KeyGen: all K, all KAT vectors
#        ./run_tb.sh top 2                # only K=2 (ML-KEM-512), all KAT vectors
#        ./run_tb.sh top 2 3              # only K=2, only CASE=3
#        ./run_tb.sh top 4 0              # only K=4 (ML-KEM-1024), CASE=0
#        ./run_tb.sh enc                  # ML-KEM Encaps: all K, all KAT vectors
#        ./run_tb.sh enc 2                # Encaps K=2, all KAT vectors
#        ./run_tb.sh enc 4 1              # Encaps K=4, only CASE=1
#        ./run_tb.sh dec                  # ML-KEM Decaps: all K, all KAT vectors
#        ./run_tb.sh dec 2 0              # Decaps K=2, only CASE=0
#        ./run_tb.sh hello [K]            # hello_world end-to-end (single instance; K=2/3/4)
#        ./run_tb.sh hello two [K]        # hello_world end-to-end (two instances; K=2/3/4)
#        ./run_tb.sh --list
#
# 'top' (KeyGen), 'enc' (Encaps) and 'dec' (Decaps) share the same RTL datapath;
# all compile the xvlog lines from sync_rtl/top/TB/xsim_run.tcl. They differ only
# in the testbench: top -> tb_mlkem_kg_katK_xsim, enc -> tb_mlkem_enc_katK_xsim,
# dec -> tb_mlkem_dec_katK_xsim.
# For these, K (2/3/4) and CASE select a subset of the KAT run so you can
# iterate quickly. Without CASE the script generates/enumerates every vector in
# test_data/kat_MLKEM_{512,768,1024}.rsp. Limit xsim parallelism with
# RUN_TB_JOBS=N (default: min(nproc, 4)).
#
# Prerequisites:
#   Vivado 2019.2 at /opt/Xilinx/Vivado/2019.2/

set -e

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
VIVADO_SETTINGS="/opt/Xilinx/Vivado/2019.2/settings64.sh"
KAT_VECTOR_DIR="$SCRIPT_DIR/sync_rtl/top/TB/vectors"
KAT_RUN_VECTOR_DIR=""

if [ "$1" = "--list" ]; then
    echo "Available modules:"
    for d in "$SCRIPT_DIR"/sync_rtl/*/TB; do
        if [ -f "$d/xsim_run.tcl" ]; then
            name=$(basename "$(dirname "$d")")
            echo "  $name"
        fi
    done
    echo "  enc   (ML-KEM Encaps; shares the 'top' RTL/tcl, enc testbench)"
    echo "  dec   (ML-KEM Decaps; shares the 'top' RTL/tcl, dec testbench)"
    echo "  hello (ML-KEM hello_world end-to-end; optional K=2/3/4, 'hello two [K]' for two instances)"
    exit 0
fi

if [ -z "$1" ]; then
    echo "Usage: $0 <module_name>"
    echo "       $0 --list"
    exit 1
fi

MODULE="$1"
SEL_K="$2"          # optional: 2/3/4 (top/enc modules only)
SEL_CASE="$3"       # optional: KAT case index (top/enc modules only)
# 'enc' is not its own RTL dir; it reuses the 'top' KeyGen tcl compile list
# (same datapath) and swaps in the encaps testbench. 'dec' likewise (Decaps).
if [ "$MODULE" = "enc" ] || [ "$MODULE" = "dec" ] || [ "$MODULE" = "hello" ]; then
    TB_DIR="$SCRIPT_DIR/sync_rtl/top/TB"
else
    TB_DIR="$SCRIPT_DIR/sync_rtl/$MODULE/TB"
fi
TCL_FILE="$TB_DIR/xsim_run.tcl"

if [ ! -f "$TCL_FILE" ]; then
    echo "ERROR: No xsim_run.tcl found for module '$MODULE'"
    echo "       Expected: $TCL_FILE"
    echo "       Run '$0 --list' to see available modules"
    exit 1
fi

# Source Vivado environment
if [ -f "$VIVADO_SETTINGS" ]; then
    source "$VIVADO_SETTINGS"
else
    echo "ERROR: Vivado settings not found at $VIVADO_SETTINGS"
    exit 1
fi

# Fix ncurses compatibility for Vivado 2019.2 on modern Linux
export LD_PRELOAD="${LD_PRELOAD}${LD_PRELOAD:+:}/usr/lib64/libtinfo.so.5"

cd "$SCRIPT_DIR"
echo "========================================"
echo " Running testbench for: $MODULE"
echo "========================================"

# Save TCL variables as bash variables, run commands with substitution
execute_tcl() {
    local tcl_file="$1"
    local tmp_script="/tmp/run_tb_$$.sh"

    # Convert TCL 'set VAR value' to bash 'VAR=value'
    {
        grep -E '^set [A-Z0-9_]+ ' "$tcl_file" | while read -r _ var value; do
            echo "${var}=${value}"
        done

        # Extract xvlog, xelab, xsim commands with bash variable expansion
        grep -E '^(xvlog|xelab|xsim) ' "$tcl_file" | while read -r cmd; do
            printf 'echo "  %s"\n' "$cmd"
            printf '%s 2>&1\n' "$cmd"
        done
    } > "$tmp_script"

    chmod +x "$tmp_script"
    . "$tmp_script"
    rm -f "$tmp_script"
}

# ---- parallel xsim launcher --------------------------------------------------
# xsim writes runtime logs into xsim.dir/<snapshot>/, so two runs of the SAME
# snapshot clobber each other. To run (K,case) in parallel we give each job its
# own COPY of xsim.dir (~1 MB, cheap) under a private workdir, run xsim there,
# and tee output to a per-job log. Compile+elaborate stay serial (they populate
# the shared xsim.dir BEFORE any parallel run starts).
#
# run_xsim_jobs <module-tag> <pass-grep-regex> <pass-match-string> <spec>...
#   each <spec> = "snapshot:K:CASE". Launches specs with bounded parallelism,
#   waits, then prints an ordered summary and returns nonzero if any job failed.
PAR_BASE="${TMPDIR:-/tmp}/run_tb_par_$$"
KAT_TMP_BASE="${TMPDIR:-/tmp}/run_tb_vectors_$$"
trap 'rm -rf "$PAR_BASE" "$KAT_TMP_BASE"' EXIT

print_timing_summary() {
    local logs=("$@")

    echo ""
    awk '
        /^TIME / {
            cycles = ""; runtime = "";
            for (i = 1; i <= NF; i++) {
                split($i, a, "=");
                if (a[1] == "cycles") cycles = a[2];
                else if (a[1] == "runtime") runtime = a[2];
            }
            if (cycles != "") {
                n++;
                sum_cycles += cycles;
                if (runtime != "") sum_runtime += runtime;
            }
        }
        END {
            if (n == 0) {
                print "Timing average: no TIME records found";
            } else {
                printf "Timing average (clk=20MHz, period=50ns): runs=%d avg_cycles=%.1f avg_time_us=%.3f avg_time_ms=%.3f\n",
                       n, sum_cycles / n, (sum_runtime / n) / 1000.0, (sum_runtime / n) / 1000000.0;
            }
        }
    ' "${logs[@]}"
}

validate_kat_selection() {
    local ksel="$1" csel="$2"
    case "$ksel" in
        ""|2|3|4) ;;
        *) echo "ERROR: K must be 2, 3, or 4 (got '$ksel')"; return 1 ;;
    esac
    if [ -n "$csel" ] && ! [[ "$csel" =~ ^[0-9]+$ ]]; then
        echo "ERROR: CASE must be a non-negative integer (got '$csel')"
        return 1
    fi
}

generate_kat_vectors() {
    local op="$1" ksel="${2:-}" csel="${3:-}"

    validate_kat_selection "$ksel" "$csel" || return 1
    KAT_RUN_VECTOR_DIR="$KAT_TMP_BASE/$op"
    rm -rf "$KAT_RUN_VECTOR_DIR"
    mkdir -p "$KAT_RUN_VECTOR_DIR"
    [ -d "$KAT_VECTOR_DIR/encgold" ] && ln -s "$KAT_VECTOR_DIR/encgold" "$KAT_RUN_VECTOR_DIR/encgold"
    [ -d "$KAT_VECTOR_DIR/decgold" ] && ln -s "$KAT_VECTOR_DIR/decgold" "$KAT_RUN_VECTOR_DIR/decgold"

    python3 - "$SCRIPT_DIR" "$op" "${ksel:-all}" "${csel:-all}" "$KAT_RUN_VECTOR_DIR" <<'PY'
import pathlib
import sys

root = pathlib.Path(sys.argv[1])
op = sys.argv[2]
ksel = sys.argv[3]
csel = sys.argv[4]
out_dir = pathlib.Path(sys.argv[5])

ks = [2, 3, 4] if ksel == "all" else [int(ksel)]
case_filter = None if csel == "all" else int(csel)
rsp_by_k = {
    2: root / "test_data" / "kat_MLKEM_512.rsp",
    3: root / "test_data" / "kat_MLKEM_768.rsp",
    4: root / "test_data" / "kat_MLKEM_1024.rsp",
}

def parse_rsp(path):
    records = []
    cur = {}
    with path.open("r", encoding="ascii") as f:
        for raw in f:
            line = raw.strip()
            if not line or line.startswith("#") or " = " not in line:
                continue
            key, val = line.split(" = ", 1)
            if key == "count":
                if cur:
                    records.append(cur)
                cur = {"count": int(val)}
            else:
                cur[key] = val.lower()
    if cur:
        records.append(cur)
    return records

def bytes_of(hexstr, field, count, k):
    if len(hexstr) % 2:
        raise SystemExit(f"odd hex length for {field}, K={k}, count={count}")
    return [hexstr[i:i + 2] for i in range(0, len(hexstr), 2)]

def write_bytes(path, hexstr, field, count, k):
    data = bytes_of(hexstr, field, count, k)
    path.write_text("".join(b + "\n" for b in data), encoding="ascii")

def write_word256_reversed(path, hexstr, field, count, k):
    data = bytes_of(hexstr, field, count, k)
    if len(data) != 32:
        raise SystemExit(f"{field} must be 32 bytes, got {len(data)}, K={k}, count={count}")
    path.write_text("".join(reversed(data)) + "\n", encoding="ascii")

def require(rec, keys, k):
    missing = [key for key in keys if key not in rec]
    if missing:
        raise SystemExit(f"missing {','.join(missing)} in K={k}, count={rec.get('count')}")

summary = []
for k in ks:
    records = parse_rsp(rsp_by_k[k])
    wrote = 0
    for rec in records:
        count = rec["count"]
        if case_filter is not None and count != case_filter:
            continue
        tag = f"k{k}"
        if op == "top":
            require(rec, ["d", "z", "pk", "sk"], k)
            base = out_dir / f"kat_{tag}_c{count}"
            write_word256_reversed(base.with_name(base.name + "_d.hex"), rec["d"], "d", count, k)
            write_word256_reversed(base.with_name(base.name + "_z.hex"), rec["z"], "z", count, k)
            write_bytes(base.with_name(base.name + "_ek.hex"), rec["pk"], "pk", count, k)
            write_bytes(base.with_name(base.name + "_dk.hex"), rec["sk"], "sk", count, k)
        elif op == "enc":
            require(rec, ["pk", "msg", "ct", "ss"], k)
            base = out_dir / f"enc_{tag}_c{count}"
            write_bytes(base.with_name(base.name + "_ek.hex"), rec["pk"], "pk", count, k)
            write_bytes(base.with_name(base.name + "_m.hex"), rec["msg"], "msg", count, k)
            write_bytes(base.with_name(base.name + "_ct.hex"), rec["ct"], "ct", count, k)
            write_bytes(base.with_name(base.name + "_ss.hex"), rec["ss"], "ss", count, k)
        elif op == "dec":
            require(rec, ["sk", "ct", "ss", "ct_n", "ss_n"], k)
            base = out_dir / f"dec_{tag}_c{count}"
            write_bytes(base.with_name(base.name + "_dk.hex"), rec["sk"], "sk", count, k)
            write_bytes(base.with_name(base.name + "_ct.hex"), rec["ct"], "ct", count, k)
            write_bytes(base.with_name(base.name + "_ss.hex"), rec["ss"], "ss", count, k)
            write_bytes(base.with_name(base.name + "_ctn.hex"), rec["ct_n"], "ct_n", count, k)
            write_bytes(base.with_name(base.name + "_ssn.hex"), rec["ss_n"], "ss_n", count, k)
        else:
            raise SystemExit(f"unknown KAT op: {op}")
        wrote += 1
    if case_filter is not None and wrote == 0:
        raise SystemExit(f"CASE={case_filter} not found in {rsp_by_k[k]}")
    summary.append(f"K={k}:{wrote}")

print(f"  generated {op} KAT vector files ({', '.join(summary)})")
PY
}

kat_cases_for() {
    local op="$1" k="$2" prefix suffix
    local dir="${KAT_RUN_VECTOR_DIR:-$KAT_VECTOR_DIR}"
    case "$op" in
        top) prefix="kat_k${k}_c"; suffix="_d\\.hex" ;;
        enc) prefix="enc_k${k}_c"; suffix="_ek\\.hex" ;;
        dec) prefix="dec_k${k}_c"; suffix="_dk\\.hex" ;;
        *) echo "ERROR: unknown KAT op '$op'" >&2; return 1 ;;
    esac

    find "$dir" -maxdepth 1 -type f -printf '%f\n' \
        | sed -n "s/^${prefix}\([0-9][0-9]*\)${suffix}$/\1/p" \
        | sort -n \
        | tr '\n' ' '
}

run_xsim_jobs() {
    local mtag="$1" pgrep="$2" pmatch="$3"; shift 3
    local specs=("$@")
    local repo; repo="$(pwd)"
    local max_jobs="${RUN_TB_JOBS:-}"
    if [ -z "$max_jobs" ]; then
        if command -v nproc >/dev/null 2>&1; then
            max_jobs="$(nproc)"
            [ "$max_jobs" -gt 4 ] && max_jobs=4
        else
            max_jobs=4
        fi
    fi
    case "$max_jobs" in
        ''|*[!0-9]*|0) echo "ERROR: RUN_TB_JOBS must be a positive integer (got '$max_jobs')"; return 1 ;;
    esac

    rm -rf "$PAR_BASE"; mkdir -p "$PAR_BASE"
    echo "  launching ${#specs[@]} xsim job(s), parallelism=$max_jobs"

    # launch all jobs in parallel
    local s snap k c wd active=0 launched=0 total_specs=${#specs[@]}
    for s in "${specs[@]}"; do
        IFS=: read -r snap k c <<< "$s"
        launched=$((launched + 1))
        if [ -t 1 ]; then
            printf "\r  running [%d/%d] K=%s CASE=%s" "$launched" "$total_specs" "$k" "$c"
        else
            echo "  running [$launched/$total_specs] K=$k CASE=$c"
        fi
        wd="$PAR_BASE/${mtag}_k${k}_c${c}"
        mkdir -p "$wd"
        cp -r "$repo/xsim.dir" "$wd/xsim.dir"
        # The TB loads vectors via repo-relative $readmemh paths. xsim resolves
        # these against the run cwd, so mount the generated vector directory
        # under each private workdir.
        if [ -n "$KAT_RUN_VECTOR_DIR" ]; then
            mkdir -p "$wd/sync_rtl/top/TB"
            ln -sf "$KAT_RUN_VECTOR_DIR" "$wd/sync_rtl/top/TB/vectors"
        else
            ln -sf "$repo/sync_rtl" "$wd/sync_rtl"
        fi
        (
            cd "$wd"
            {
                echo "xsim $snap -R -testplusarg CASE=$c"
                xsim "$snap" -R -testplusarg "CASE=$c" --nolog
            } > "$wd/run.out" 2>&1
        ) &
        active=$((active + 1))
        if [ "$active" -ge "$max_jobs" ]; then
            wait -n || true
            active=$((active - 1))
        fi
    done
    wait || true
    [ -t 1 ] && echo ""

    # ordered result aggregation + per-job log copy to /tmp/run_tb_<mtag>_k*_c*.log
    local fail=0 pf nf log ok
    local total=0 pass_count=0 fail_count=0 noresult_count=0 fnf_count=0
    local timing_logs=() failures=()
    for s in "${specs[@]}"; do
        IFS=: read -r snap k c <<< "$s"
        wd="$PAR_BASE/${mtag}_k${k}_c${c}"
        log="/tmp/run_tb_${mtag}_k${k}_c${c}.log"
        cp -f "$wd/run.out" "$log" 2>/dev/null
        timing_logs+=("$log")
        pf=$(grep -oE "$pgrep" "$log" | tail -1)
        nf=$(grep -c 'cannot be opened' "$log")
        total=$((total + 1))
        [ -n "$pf" ] || noresult_count=$((noresult_count + 1))
        [ "$nf" -eq 0 ] || fnf_count=$((fnf_count + 1))

        # pmatch == "PASS*" means "any PASS (Dn) stage tag"; else exact match.
        if [ "$pmatch" = "PASS*" ]; then
            [[ "$pf" == PASS* ]] && [ "$nf" -eq 0 ]; ok=$?
        else
            [ "$pf" = "$pmatch" ] && [ "$nf" -eq 0 ]; ok=$?
        fi

        if [ "$ok" -eq 0 ]; then
            pass_count=$((pass_count + 1))
        else
            fail=1
            fail_count=$((fail_count + 1))
            if [ "${#failures[@]}" -lt 20 ]; then
                failures+=("K=$k CASE=$c -> ${pf:-NORESULT} (file-not-found=$nf, log: $log)")
            fi
        fi
    done

    echo ""
    echo "Result summary: total=$total success=$pass_count failure=$fail_count noresult=$noresult_count file-not-found=$fnf_count"
    if [ "$fail_count" -gt 0 ]; then
        echo "Failed cases (first ${#failures[@]} of $fail_count):"
        local item
        for item in "${failures[@]}"; do
            echo "    $item"
        done
    fi

    print_timing_summary "${timing_logs[@]}"
    rm -rf "$PAR_BASE"
    return $fail
}

# Fast single-K / single-case path for the 'top' (ML-KEM KeyGen) module.
# Compiles the same RTL as the tcl, but elaborates ONLY the requested K and
# runs the generated/enumerated CASE(s) through the bounded parallel launcher.
run_top_selected() {
    local tcl_file="$1" ksel="$2" csel="$3"
    set +e   # we do our own error handling / result parsing here

    generate_kat_vectors top "$ksel" "$csel" || return 1
    rm -rf xsim.dir .Xil

    # Compile every xvlog line from the tcl, verbatim.
    while read -r cmd; do
        eval "$cmd" || { echo "COMPILE FAILED: $cmd"; return 1; }
    done < <(grep -E '^xvlog ' "$tcl_file")

    # K to run: requested one, or all of 2/3/4. Elaborate snapshots (serial).
    local ks; if [ -n "$ksel" ]; then ks="$ksel"; else ks="2 3 4"; fi
    local k specs=()
    for k in $ks; do
        echo "  xelab tb_mlkem_kg_katK_xsim -generic_top KP=$k -s mlkem_kg_k$k --timescale 1ns/1ps"
        xelab tb_mlkem_kg_katK_xsim -generic_top KP=$k -s mlkem_kg_k$k --timescale 1ns/1ps \
            || { echo "ELAB FAILED for K=$k"; return 1; }
        local cases
        if [ -n "$csel" ]; then
            cases="$csel"
        else
            cases="$(kat_cases_for top "$k")"
        fi
        [ -n "$cases" ] || { echo "ERROR: no KeyGen KAT cases found for K=$k"; return 1; }
        local c
        for c in $cases; do specs+=("mlkem_kg_k$k:$k:$c"); done
    done

    run_xsim_jobs top 'PASS|FAIL' PASS "${specs[@]}"
}

# ML-KEM Encaps runner. Compiles the 'top' tcl xvlog lines (KeyGen datapath =
# Encaps datapath) MINUS the KeyGen TB, plus the encaps TB, then elaborates the
# requested K and runs the requested CASE(s). Mirrors run_top_selected.
run_enc_selected() {
    local tcl_file="$1" ksel="$2" csel="$3"
    set +e

    generate_kat_vectors enc "$ksel" "$csel" || return 1
    rm -rf xsim.dir .Xil

    # Compile every xvlog line from the tcl EXCEPT the KeyGen TB (the encaps TB
    # replaces it). The tcl already includes pipeline_reg + comp_decomp_sync.
    while read -r cmd; do
        [[ "$cmd" == *tb_mlkem_kg_katK* ]] && continue
        eval "$cmd" || { echo "COMPILE FAILED: $cmd"; return 1; }
    done < <(grep -E '^xvlog ' "$tcl_file")

    echo "  xvlog -sv --relax sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v"
    xvlog -sv --relax sync_rtl/top/TB/tb_mlkem_enc_katK_xsim.v \
        || { echo "ENC TB COMPILE FAILED"; return 1; }

    # K to run: requested one, or all of 2/3/4. Elaborate snapshots (serial),
    # then run all (K,case) in parallel.
    local ks; if [ -n "$ksel" ]; then ks="$ksel"; else ks="2 3 4"; fi
    local k specs=()
    for k in $ks; do
        echo "  xelab tb_mlkem_enc_katK_xsim -generic_top KP=$k -s mlkem_enc_k$k --timescale 1ns/1ps"
        xelab tb_mlkem_enc_katK_xsim -generic_top KP=$k -s mlkem_enc_k$k --timescale 1ns/1ps \
            || { echo "ELAB FAILED for K=$k"; return 1; }
        local cases
        if [ -n "$csel" ]; then
            cases="$csel"
        else
            cases="$(kat_cases_for enc "$k")"
        fi
        [ -n "$cases" ] || { echo "ERROR: no Encaps KAT cases found for K=$k"; return 1; }
        local c
        for c in $cases; do specs+=("mlkem_enc_k$k:$k:$c"); done
    done

    run_xsim_jobs enc 'PASS \(E7\)|FAIL \(E7\)' 'PASS (E7)' "${specs[@]}"
}

# ML-KEM Decaps runner. Same compile basis as enc (top tcl minus KeyGen TB) plus
# the decaps TB. Result line matches the final-stage PASS tag (D0..D7).
run_dec_selected() {
    local tcl_file="$1" ksel="$2" csel="$3"
    set +e

    generate_kat_vectors dec "$ksel" "$csel" || return 1
    rm -rf xsim.dir .Xil

    while read -r cmd; do
        [[ "$cmd" == *tb_mlkem_kg_katK* ]] && continue
        eval "$cmd" || { echo "COMPILE FAILED: $cmd"; return 1; }
    done < <(grep -E '^xvlog ' "$tcl_file")

    echo "  xvlog -sv --relax sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v"
    xvlog -sv --relax sync_rtl/top/TB/tb_mlkem_dec_katK_xsim.v \
        || { echo "DEC TB COMPILE FAILED"; return 1; }

    local ks; if [ -n "$ksel" ]; then ks="$ksel"; else ks="2 3 4"; fi
    local k specs=()
    for k in $ks; do
        echo "  xelab tb_mlkem_dec_katK_xsim -generic_top KP=$k -s mlkem_dec_k$k --timescale 1ns/1ps"
        xelab tb_mlkem_dec_katK_xsim -generic_top KP=$k -s mlkem_dec_k$k --timescale 1ns/1ps \
            || { echo "ELAB FAILED for K=$k"; return 1; }
        local cases
        if [ -n "$csel" ]; then
            cases="$csel"
        else
            cases="$(kat_cases_for dec "$k")"
        fi
        [ -n "$cases" ] || { echo "ERROR: no Decaps KAT cases found for K=$k"; return 1; }
        local c
        for c in $cases; do specs+=("mlkem_dec_k$k:$k:$c"); done
    done

    # match the highest-stage result line: 'K=.. CASE .. PASS (Dn): ...'
    run_xsim_jobs dec 'PASS \(D[0-7]\)|FAIL \(D[0-7]\)' 'PASS*' "${specs[@]}"
}

if [ "$MODULE" = "enc" ]; then
    run_enc_selected "$TCL_FILE" "$SEL_K" "$SEL_CASE"
    exit $?
fi

if [ "$MODULE" = "dec" ]; then
    run_dec_selected "$TCL_FILE" "$SEL_K" "$SEL_CASE"
    exit $?
fi

if [ "$MODULE" = "top" ]; then
    # run_top_selected handles all-K (no SEL_K) or a single K, cases in parallel.
    run_top_selected "$TCL_FILE" "$SEL_K" "$SEL_CASE"
    exit $?
fi

# ML-KEM hello_world end-to-end. Reuses the 'top' RTL compile list, then runs a
# self-checking protocol TB (no KAT vectors): single instance by default, or two
# instances (genenc + dec) with 'hello two'. Optional K selects ML-KEM-512/768/1024.
run_hello_selected() {
    local tcl_file="$1" arg1="$2" arg2="$3"
    set +e
    rm -rf xsim.dir .Xil

    local variant="" ksel="2"
    if [ "$arg1" = "two" ]; then
        variant="two"
        [ -n "$arg2" ] && ksel="$arg2"
    elif [ -n "$arg1" ]; then
        ksel="$arg1"
    fi
    case "$ksel" in
        2|3|4) ;;
        *) echo "ERROR: hello K must be 2, 3, or 4 (got '$ksel')"; return 1 ;;
    esac

    local tb snap
    if [ "$variant" = "two" ]; then
        tb=tb_mlkem_two_inst_xsim;  snap=mlkem_two
    else
        tb=tb_mlkem_hello_world_xsim; snap=mlkem_hello
    fi

    # compile RTL (every non-TB xvlog line from the shared tcl)
    while read -r cmd; do
        [[ "$cmd" == *"TB/tb_"* ]] && continue
        eval "$cmd" || { echo "COMPILE FAILED: $cmd"; return 1; }
    done < <(grep -E '^xvlog ' "$tcl_file")

    echo "  xvlog -sv --relax sync_rtl/top/TB/$tb.v"
    xvlog -sv --relax "sync_rtl/top/TB/$tb.v" || { echo "HELLO TB COMPILE FAILED"; return 1; }

    echo "  xelab $tb -generic_top KP=$ksel -s ${snap}_k${ksel} --timescale 1ns/1ps"
    xelab "$tb" -generic_top "KP=$ksel" -s "${snap}_k${ksel}" --timescale 1ns/1ps || { echo "ELAB FAILED"; return 1; }

    xsim "${snap}_k${ksel}" -R
}

if [ "$MODULE" = "hello" ]; then
    run_hello_selected "$TCL_FILE" "$SEL_K" "$SEL_CASE"  # args are [K] or [two] [K]
    exit $?
fi

execute_tcl "$TCL_FILE"