#!/bin/bash
# run_tb.sh - Run Vivado XSIM testbench for a module
#
# Usage: ./run_tb.sh <module> [K] [CASE]
#        ./run_tb.sh mod_add              # run the module's full xsim_run.tcl
#        ./run_tb.sh top                  # ML-KEM KeyGen: all K, all cases
#        ./run_tb.sh top 2                # only K=2 (ML-KEM-512), all its cases
#        ./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 --list
#
# For the 'top' module, K (2/3/4) and CASE select a single KAT run so you can
# iterate quickly instead of elaborating 3 snapshots and running all 11 cases.
# Other modules ignore the extra args and run their xsim_run.tcl verbatim.
#
# 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"

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
    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 module only)
SEL_CASE="$3"       # optional: KAT case index (top module only)
TB_DIR="$SCRIPT_DIR/sync_rtl/$MODULE/TB"
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"
}

# 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 ONLY the requested CASE(s). Avoids the 3-snapshot + 11-run full sweep.
run_top_selected() {
    local tcl_file="$1" k="$2" csel="$3"
    set +e   # we do our own error handling / result parsing here

    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")

    # Cases per K: K=2 has 0..4, K=3/4 have 0..2.
    local cases
    if [ "$k" = "2" ]; then cases="0 1 2 3 4"; else cases="0 1 2"; fi
    if [ -n "$csel" ]; then cases="$csel"; fi

    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 fail=0
    for c in $cases; do
        local log="/tmp/run_tb_top_k${k}_c${c}.log"
        echo "  xsim mlkem_kg_k$k -R -testplusarg CASE=$c   (log: $log)"
        xsim "mlkem_kg_k$k" -R -testplusarg "CASE=$c" > "$log" 2>&1
        local pf nf
        pf=$(grep -oE 'PASS|FAIL' "$log" | tail -1)
        nf=$(grep -c 'cannot be opened' "$log")
        echo "    K=$k CASE=$c -> ${pf:-NORESULT}  (file-not-found=$nf)"
        { [ "$pf" = "PASS" ] && [ "$nf" -eq 0 ]; } || fail=1
    done
    return $fail
}

if [ "$MODULE" = "top" ] && [ -n "$SEL_K" ]; then
    run_top_selected "$TCL_FILE" "$SEL_K" "$SEL_CASE"
    exit $?
fi

execute_tcl "$TCL_FILE"
