#!/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 enc                  # ML-KEM Encaps: all K, all cases (0..2)
#        ./run_tb.sh enc 2                # Encaps K=2, all its cases
#        ./run_tb.sh enc 4 1              # Encaps K=4, only CASE=1
#        ./run_tb.sh --list
#
# 'top' (KeyGen) and 'enc' (Encaps) share the same RTL datapath; both 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.
# For these two, K (2/3/4) and CASE select a single KAT run so you can iterate
# quickly. 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
    echo "  enc   (ML-KEM Encaps; shares the 'top' RTL/tcl, enc testbench)"
    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.
if [ "$MODULE" = "enc" ]; 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"
}

# 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"
        echo "========================================" | tee "$log"
        xsim "mlkem_kg_k$k" -R -testplusarg "CASE=$c" 2>&1 | tee -a "$log"
        echo "========================================" | tee -a "$log"
        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, log: $log)"
        { [ "$pf" = "PASS" ] && [ "$nf" -eq 0 ]; } || fail=1
    done
    return $fail
}

# 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

    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.
    local ks; if [ -n "$ksel" ]; then ks="$ksel"; else ks="2 3 4"; fi
    local fail=0
    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"; fail=1; continue; }

        # Encaps KAT cases per K: 0..2 (vectors enc_k{K}_c{0,1,2}_*).
        local cases="0 1 2"
        if [ -n "$csel" ]; then cases="$csel"; fi
        for c in $cases; do
            local log="/tmp/run_tb_enc_k${k}_c${c}.log"
            echo "  xsim mlkem_enc_k$k -R -testplusarg CASE=$c"
            echo "========================================" | tee "$log"
            xsim "mlkem_enc_k$k" -R -testplusarg "CASE=$c" 2>&1 | tee -a "$log"
            echo "========================================" | tee -a "$log"
            local pf nf
            pf=$(grep -oE 'PASS \(E7\)|FAIL \(E7\)' "$log" | tail -1)
            nf=$(grep -c 'cannot be opened' "$log")
            echo "    K=$k CASE=$c -> ${pf:-NORESULT}  (file-not-found=$nf, log: $log)"
            { [ "$pf" = "PASS (E7)" ] && [ "$nf" -eq 0 ]; } || fail=1
        done
    done
    return $fail
}

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

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

execute_tcl "$TCL_FILE"