#!/usr/bin/env bash

set -euo pipefail

###################################
# Defaults / Configuration
###################################
default_vasp="vasp_std"
default_num_openmp="${OMP_NUM_THREADS:-1}"
default_num_mkl="${MKL_NUM_THREADS:-1}"
ownership="tkamiya.csrv"

###################################
# Helper functions
###################################
die() {
    printf 'ERROR: %s\n' "$*" >&2
    exit 1
}

usage() {
    cat <<EOF_USAGE
Usage:
  $0 [vasp_path] [n_mpirun_processors] [num_openmp] [num_mkl]

Arguments:
  vasp_path               VASP executable name or full/relative path
                          default: ${default_vasp}
  n_mpirun_processors     Number of MPI processes
                          default: number of physical CPU cores
  num_openmp              OMP_NUM_THREADS
                          default: current OMP_NUM_THREADS, or 1 if unset
  num_mkl                 MKL_NUM_THREADS
                          default: current MKL_NUM_THREADS, or 1 if unset

Examples:
  $0
  $0 vasp_std
  $0 vasp_std 4 1 1
  $0 /home/share/vasp.6.6.1/bin/vasp_std 4 1 1
EOF_USAGE
}

is_positive_integer() {
    [[ $1 =~ ^[1-9][0-9]*$ ]]
}

resolve_executable() {
    local spec=$1
    local resolved

    if [[ $spec == */* ]]; then
        [[ -f $spec ]] || die "VASP executable not found: $spec"
        [[ -x $spec ]] || die "VASP file is not executable: $spec"
        resolved=$(readlink -f -- "$spec")
    else
        resolved=$(command -v -- "$spec" 2>/dev/null) || \
            die "VASP executable not found in PATH: $spec"
    fi

    printf '%s\n' "$resolved"
}

escape_sed_replacement() {
    # Escape characters special in the replacement part of sed s|||.
    printf '%s' "$1" | sed 's/[\\&|]/\\&/g'
}

###################################
# Host information
###################################
command -v lscpu >/dev/null 2>&1 || die "lscpu command not found"
command -v nproc >/dev/null 2>&1 || die "nproc command not found"

nSockets_host=$(
    lscpu -p=SOCKET |
        grep -v '^#' |
        sort -u |
        wc -l
)

nLogicalProcessors_host=$(nproc)

nThreadsPerCore_host=$(
    lscpu |
        awk -F: '/Thread\(s\) per core/ {
            gsub(/^[ \t]+/, "", $2)
            print $2
            exit
        }'
)

# Some ARM/WSL2 environments expose unreliable or duplicated CORE IDs via
# `lscpu -p=SOCKET,CORE`.  Derive the physical core count from the number of
# online logical processors and threads per core instead.  This also works for
# conventional SMT systems (e.g. 4 cores / 8 threads => 8 / 2 = 4).
if [[ $nThreadsPerCore_host =~ ^[1-9][0-9]*$ ]] &&
   (( nLogicalProcessors_host % nThreadsPerCore_host == 0 )); then
    nPhysicalCores_host=$(( nLogicalProcessors_host / nThreadsPerCore_host ))
else
    # Fallback for unusual lscpu output.
    nPhysicalCores_host=$(
        lscpu -p=SOCKET,CORE |
            grep -v '^#' |
            sort -u |
            wc -l
    )
fi

CPUModel_host=$(
    lscpu |
        awk -F: '
            /Model name/ {
                gsub(/^[ \t]+/, "", $2)
                if ($2 != "") {
                    print $2
                    exit
                }
            }
            /^Model:/ {
                gsub(/^[ \t]+/, "", $2)
                if ($2 != "") {
                    print $2
                    exit
                }
            }
        '
)

if [[ -z $CPUModel_host ]]; then
    CPUModel_host=$(
        awk -F: '
            /model name/ || /Processor/ || /Hardware/ {
                gsub(/^[ \t]+/, "", $2)
                if ($2 != "") {
                    print $2
                    exit
                }
            }
        ' /proc/cpuinfo
    )
fi

CPUModel_host=${CPUModel_host:-unknown}

# OS and execution-environment information.  /etc/os-release is the standard
# source for Linux distribution names and versions.  WSL 2 can be identified
# from the Microsoft WSL2 kernel release; an older Microsoft kernel denotes
# WSL 1.
OSName_host=unknown
OSVersion_host=unknown
OSPrettyName_host=unknown

if [[ -r /etc/os-release ]]; then
    OSName_host=$(
        . /etc/os-release
        printf '%s\n' "${NAME:-unknown}"
    )
    OSVersion_host=$(
        . /etc/os-release
        printf '%s\n' "${VERSION_ID:-${VERSION:-unknown}}"
    )
    OSPrettyName_host=$(
        . /etc/os-release
        printf '%s\n' "${PRETTY_NAME:-${NAME:-unknown} ${VERSION_ID:-}}"
    )
fi

KernelRelease_host=$(uname -r)
Architecture_host=$(uname -m)
kernel_release_lower=${KernelRelease_host,,}

if [[ $kernel_release_lower == *microsoft* || -n ${WSL_INTEROP:-} ]]; then
    if [[ $kernel_release_lower == *wsl2* ||
          $kernel_release_lower == *microsoft-standard* ]]; then
        ExecutionEnvironment_host="WSL 2"
    else
        ExecutionEnvironment_host="WSL 1"
    fi
else
    ExecutionEnvironment_host="native Linux"
fi

###################################
# Command-line arguments
###################################
if (( $# > 4 )); then
    usage >&2
    exit 2
fi

if [[ ${1:-} == "-h" || ${1:-} == "--help" ]]; then
    usage
    exit 0
fi

vasp_spec=${1:-$default_vasp}
nMPI=${2:-$nPhysicalCores_host}
num_openmp=${3:-$default_num_openmp}
num_mkl=${4:-$default_num_mkl}

is_positive_integer "$nMPI"       || die "n_mpirun_processors must be a positive integer: $nMPI"
is_positive_integer "$num_openmp" || die "num_openmp must be a positive integer: $num_openmp"
is_positive_integer "$num_mkl"    || die "num_mkl must be a positive integer: $num_mkl"

vasp_path=$(resolve_executable "$vasp_spec")
vasp_dir=$(dirname -- "$vasp_path")

###################################
# Paths / Working parameters
###################################
root_dir=$PWD
template_dir="$root_dir/template"
input_template_dir="$template_dir/test_inputs"
run_sh_template_path="$template_dir/run_test_template.sh"

hostname=$(hostname)
date=$(date '+%Y%m%d')
work_dir="$root_dir/${hostname}_np${nMPI}_omp${num_openmp}_mkl${num_mkl}_${date}"

md_path="$work_dir/condition.md"
run_sh_path="$work_dir/run.sh"

[[ -d $input_template_dir ]] || die "Input template directory not found: $input_template_dir"
[[ -f $run_sh_template_path ]] || die "Run-script template not found: $run_sh_template_path"
[[ ! -e $work_dir ]] || die "Work directory already exists: $work_dir"

# %q makes each value safe as a Bash word in the generated command.
printf -v vasp_command \
    'env OMP_NUM_THREADS=%q MKL_NUM_THREADS=%q mpirun -np %q %q 2>&1 | tee vasp.log' \
    "$num_openmp" "$num_mkl" "$nMPI" "$vasp_path"

escaped_work_dir=$(escape_sed_replacement "$work_dir")
escaped_vasp_command=$(escape_sed_replacement "$vasp_command")

###################################
# Initialization
###################################
printf '\n'
printf 'root_dir           : %s\n' "$root_dir"
printf 'template_dir       : %s\n' "$template_dir"
printf 'input_template_dir : %s\n' "$input_template_dir"
printf '\n'
printf 'hostname           : %s\n' "$hostname"
printf 'account            : %s\n' "$(whoami)"
printf 'date               : %s\n' "$date"
printf '\n'
printf 'OS                 : %s\n' "$OSPrettyName_host"
printf 'OS name            : %s\n' "$OSName_host"
printf 'OS version         : %s\n' "$OSVersion_host"
printf 'Execution env.     : %s\n' "$ExecutionEnvironment_host"
printf 'Kernel release     : %s\n' "$KernelRelease_host"
printf 'Architecture       : %s\n' "$Architecture_host"
printf '\n'
printf 'CPU model          : %s\n' "$CPUModel_host"
printf 'Sockets            : %s\n' "$nSockets_host"
printf 'Physical cores     : %s\n' "$nPhysicalCores_host"
printf 'Threads/core       : %s\n' "$nThreadsPerCore_host"
printf 'Logical processors : %s\n' "$nLogicalProcessors_host"
printf 'MPI processes      : %s\n' "$nMPI"
printf 'OMP_NUM_THREADS    : %s\n' "$num_openmp"
printf 'MKL_NUM_THREADS    : %s\n' "$num_mkl"
printf '\n'
printf 'work_dir           : %s\n' "$work_dir"
printf 'vasp_path          : %s\n' "$vasp_path"
printf 'vasp_dir           : %s\n' "$vasp_dir"
printf 'run_sh_path        : %s\n' "$run_sh_path"

printf '\nCreating work directory: %s\n' "$work_dir"
mkdir -- "$work_dir"

printf 'Copying test inputs...\n'
cp -a "$input_template_dir"/. "$work_dir"/

# Copy VASP makefiles when present.  Absence is not fatal because an installed
# executable may live separately from the build tree.
shopt -s nullglob
makefiles=("$vasp_dir"/makefile*)
shopt -u nullglob

if (( ${#makefiles[@]} > 0 )); then
    printf 'Copying VASP makefile(s)...\n'
    cp -a -- "${makefiles[@]}" "$work_dir"/
else
    printf 'Note: no makefile* found in %s; skipping makefile copy.\n' "$vasp_dir"
fi

###################################
# condition.md
###################################
printf 'Creating %s\n' "$md_path"

cat > "$md_path" <<EOF_MD
# Configuration
hostname: $hostname
account: $(whoami)
date: $(date)
work_dir: $work_dir
vasp_path: $vasp_path
vasp_dir: $vasp_dir

# Machine information
OS: $OSPrettyName_host
OS name: $OSName_host
OS version: $OSVersion_host
Execution environment: $ExecutionEnvironment_host
Kernel release: $KernelRelease_host
Architecture: $Architecture_host
CPU model: $CPUModel_host
Sockets: $nSockets_host
Physical cores: $nPhysicalCores_host
Threads per core: $nThreadsPerCore_host
Logical processors: $nLogicalProcessors_host
MPI processes: $nMPI
OMP_NUM_THREADS: $num_openmp
MKL_NUM_THREADS: $num_mkl

# VASP command
$vasp_command

# VASP linked libraries
~~~text
EOF_MD

if command -v ldd >/dev/null 2>&1; then
    printf '$ ldd %q\n' "$vasp_path" >> "$md_path"
    ldd_status=0
    ldd "$vasp_path" >> "$md_path" 2>&1 || ldd_status=$?
    if (( ldd_status != 0 )); then
        printf '[ldd exited with status %d]\n' "$ldd_status" >> "$md_path"
    fi
else
    printf 'ldd command not found; linked libraries could not be recorded.\n' >> "$md_path"
fi

cat >> "$md_path" <<EOF_MD
~~~

# Environment
EOF_MD

env | sort >> "$md_path"

###################################
# Generate benchmark run script
###################################
printf 'Creating %s\n' "$run_sh_path"

sed \
    -e "s|{{vasp_command}}|${escaped_vasp_command}|g" \
    -e "s|{{work_dir}}|${escaped_work_dir}|g" \
    "$run_sh_template_path" > "$run_sh_path"

chmod +x "$run_sh_path"

printf 'Run test script [%s] has been created.\n' "$run_sh_path"

###################################
# Termination process
###################################
if (( EUID == 0 )); then
    printf '\nChange ownership/permission from root to %s:770\n' "$ownership"
    chown -R "$ownership" "$work_dir"
    chmod -R 770 "$work_dir"
fi

printf '\nDone.\n'