#!/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 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 dec                  # ML-KEM Decaps: all K, all cases (0..2)
#        ./run_tb.sh dec 2 0              # Decaps K=2, only CASE=0
#        ./run_tb.sh hello                # hello_world end-to-end (single instance)
#        ./run_tb.sh hello two            # hello_world end-to-end (two instances)
#        ./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 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)"
    echo "  dec   (ML-KEM Decaps; shares the 'top' RTL/tcl, dec testbench)"
    echo "  hello (ML-KEM hello_world end-to-end; 'hello two' 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 all specs in parallel, waits,
#   then prints an ordered summary and returns nonzero if any job failed.
PAR_BASE="${TMPDIR:-/tmp}/run_tb_par_$$"
run_xsim_jobs() {
    local mtag="$1" pgrep="$2" pmatch="$3"; shift 3
    local specs=("$@")
    local repo; repo="$(pwd)"
    rm -rf "$PAR_BASE"; mkdir -p "$PAR_BASE"

    # launch all jobs in parallel
    local s snap k c wd
    for s in "${specs[@]}"; do
        IFS=: read -r snap k c <<< "$s"
        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
        # (sync_rtl/top/TB/vectors/...). xsim resolves these against the run
        # cwd, so symlink the repo's sync_rtl tree into each private workdir.
        ln -sf "$repo/sync_rtl" "$wd/sync_rtl"
        (
            cd "$wd"
            echo "  xsim $snap -R -testplusarg CASE=$c"
            xsim "$snap" -R -testplusarg "CASE=$c" --nolog > "$wd/run.out" 2>&1
        ) &
    done
    wait

    # ordered summary + per-job log copy to /tmp/run_tb_<mtag>_k*_c*.log
    local fail=0 pf nf log
    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
        pf=$(grep -oE "$pgrep" "$log" | tail -1)
        nf=$(grep -c 'cannot be opened' "$log")
        echo "    K=$k CASE=$c -> ${pf:-NORESULT}  (file-not-found=$nf, log: $log)"
        # pmatch == "PASS*" means "any PASS (Dn) stage tag"; else exact match.
        if [ "$pmatch" = "PASS*" ]; then
            { [[ "$pf" == PASS* ]] && [ "$nf" -eq 0 ]; } || fail=1
        else
            { [ "$pf" = "$pmatch" ] && [ "$nf" -eq 0 ]; } || fail=1
        fi
    done
    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 CASE(s) in parallel. Avoids the 3-snapshot + 11-run serial sweep.
run_top_selected() {
    local tcl_file="$1" ksel="$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")

    # 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; }
        # 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
        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

    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="0 1 2"
        if [ -n "$csel" ]; then cases="$csel"; fi
        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

    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="0 1 2"
        if [ -n "$csel" ]; then cases="$csel"; fi
        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'. Prints each step's inputs/outputs.
run_hello_selected() {
    local tcl_file="$1" variant="$2"
    set +e
    rm -rf xsim.dir .Xil

    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 -s $snap --timescale 1ns/1ps"
    xelab "$tb" -s "$snap" --timescale 1ns/1ps || { echo "ELAB FAILED"; return 1; }

    xsim "$snap" -R
}

if [ "$MODULE" = "hello" ]; then
    run_hello_selected "$TCL_FILE" "$SEL_K"   # SEL_K position carries the variant ('two')
    exit $?
fi

execute_tcl "$TCL_FILE"