#!/bin/bash
# run_enc.sh - Run Vivado XSIM for the ML-KEM Encaps testbench (tb_mlkem_enc_katK).
#
# Usage: ./run_enc.sh [K] [CASE]
#        ./run_enc.sh          # K=2 CASE=0
#        ./run_enc.sh 3 1      # K=3 CASE=1
#
# Compiles the same RTL as the KeyGen tcl (mlkem_top + leaves), elaborates the
# encaps TB for the requested K, and runs the requested CASE.

set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
VIVADO_SETTINGS="/opt/Xilinx/Vivado/2019.2/settings64.sh"

SEL_K="${1:-2}"
SEL_CASE="${2:-0}"

source "$VIVADO_SETTINGS"
export LD_PRELOAD="${LD_PRELOAD}${LD_PRELOAD:+:}/usr/lib64/libtinfo.so.5"

cd "$SCRIPT_DIR"
rm -rf xsim.dir .Xil

# Compile every xvlog RTL line from the KeyGen tcl (same datapath), verbatim.
TCL="sync_rtl/top/TB/xsim_run.tcl"
while read -r cmd; do
    [[ "$cmd" == *tb_mlkem_kg_katK* ]] && continue   # skip the KeyGen TB
    [[ "$cmd" == *mlkem_top.v* ]] && continue        # defer top until comp_decomp compiled
    echo "  $cmd"
    eval "$cmd" || { echo "COMPILE FAILED: $cmd"; exit 1; }
done < <(grep -E '^xvlog ' "$TCL")

# Encaps-only leaf: comp_decomp_sync (compress) + its pipeline_reg dep.
for f in sync_rtl/common/pipeline_reg.v sync_rtl/comp_decomp/comp_decomp_sync.v; do
    [[ -f "$f" ]] || continue
    echo "  xvlog -sv --relax -i . $f"
    xvlog -sv --relax -i . "$f" || { echo "COMPILE FAILED: $f"; exit 1; }
done

# Now compile the top (depends on comp_decomp_sync).
echo "  xvlog -sv --relax -i . sync_rtl/top/mlkem_top.v"
xvlog -sv --relax -i . sync_rtl/top/mlkem_top.v \
    || { echo "TOP COMPILE FAILED"; exit 1; }

# Compile the encaps TB
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 "TB COMPILE FAILED"; exit 1; }

echo "  xelab tb_mlkem_enc_katK_xsim -generic_top KP=$SEL_K -s mlkem_enc_k$SEL_K --timescale 1ns/1ps"
xelab tb_mlkem_enc_katK_xsim -generic_top KP=$SEL_K -s mlkem_enc_k$SEL_K --timescale 1ns/1ps \
    || { echo "ELAB FAILED for K=$SEL_K"; exit 1; }

LOG="/tmp/run_enc_k${SEL_K}_c${SEL_CASE}.log"
echo "  xsim mlkem_enc_k$SEL_K -R -testplusarg CASE=$SEL_CASE"
xsim "mlkem_enc_k$SEL_K" -R -testplusarg "CASE=$SEL_CASE" 2>&1 | tee "$LOG"