#!/bin/sh
set -eu

DURATION_SEC=300
DEVICES=""
EXCLUDE=""
FORMAT=""
TEST_SLICE_SECONDS=300
JOHN_DIR="/usr/local/lib/bee/john/run"
JOHN_BIN="${JOHN_DIR}/john"
export OCL_ICD_VENDORS="/etc/OpenCL/vendors"
export LD_LIBRARY_PATH="/usr/lib:/usr/local/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}"

usage() {
    echo "usage: $0 [--seconds N] [--devices 0,1] [--exclude 2,3] [--format name]" >&2
    exit 2
}

normalize_list() {
    echo "${1:-}" | tr ',' '\n' | sed 's/[[:space:]]//g' | awk 'NF' | sort -n | uniq | paste -sd, -
}

contains_csv() {
    needle="$1"
    haystack="${2:-}"
    echo ",${haystack}," | grep -q ",${needle},"
}

show_opencl_diagnostics() {
    echo "-- OpenCL ICD vendors --" >&2
    if [ -d /etc/OpenCL/vendors ]; then
        ls -l /etc/OpenCL/vendors >&2 || true
        for icd in /etc/OpenCL/vendors/*.icd; do
            [ -f "${icd}" ] || continue
            echo "  file: ${icd}" >&2
            sed 's/^/    /' "${icd}" >&2 || true
        done
    else
        echo "  /etc/OpenCL/vendors is missing" >&2
    fi
    echo "-- NVIDIA device nodes --" >&2
    ls -l /dev/nvidia* >&2 || true
    echo "-- ldconfig OpenCL/NVIDIA --" >&2
    ldconfig -p 2>/dev/null | grep 'libOpenCL\|libcuda\|libnvidia-opencl' >&2 || true
    if command -v clinfo >/dev/null 2>&1; then
        echo "-- clinfo -l --" >&2
        clinfo -l >&2 || true
    fi
    echo "-- john --list=opencl-devices --" >&2
    ./john --list=opencl-devices >&2 || true
}

refresh_nvidia_runtime() {
    if [ "$(id -u)" != "0" ]; then
        return 1
    fi
    if command -v bee-nvidia-load >/dev/null 2>&1; then
        bee-nvidia-load >/dev/null 2>&1 || true
    fi
    ldconfig >/dev/null 2>&1 || true
    return 0
}

ensure_nvidia_uvm() {
    if lsmod 2>/dev/null | grep -q '^nvidia_uvm '; then
        return 0
    fi
    if [ "$(id -u)" != "0" ]; then
        return 1
    fi

    ko="/usr/local/lib/nvidia/nvidia-uvm.ko"
    [ -f "${ko}" ] || return 1

    if ! insmod "${ko}" >/dev/null 2>&1; then
        return 1
    fi

    uvm_major=$(grep -m1 ' nvidia-uvm$' /proc/devices | awk '{print $1}')
    if [ -n "${uvm_major}" ]; then
        mknod -m 666 /dev/nvidia-uvm c "${uvm_major}" 0 2>/dev/null || true
        mknod -m 666 /dev/nvidia-uvm-tools c "${uvm_major}" 1 2>/dev/null || true
    fi
    return 0
}

ensure_opencl_ready() {
    out=$(./john --list=opencl-devices 2>&1 || true)
    if echo "${out}" | grep -q "Device #"; then
        return 0
    fi

    if refresh_nvidia_runtime; then
        out=$(./john --list=opencl-devices 2>&1 || true)
        if echo "${out}" | grep -q "Device #"; then
            return 0
        fi
    fi

    if ensure_nvidia_uvm; then
        out=$(./john --list=opencl-devices 2>&1 || true)
        if echo "${out}" | grep -q "Device #"; then
            return 0
        fi
    fi

    echo "OpenCL devices are not available for John." >&2
    if ! lsmod 2>/dev/null | grep -q '^nvidia_uvm '; then
        echo "nvidia_uvm is not loaded." >&2
    fi
    if [ ! -e /dev/nvidia-uvm ]; then
        echo "/dev/nvidia-uvm is missing." >&2
    fi
    show_opencl_diagnostics
    return 1
}

while [ "$#" -gt 0 ]; do
    case "$1" in
        --seconds|-t) [ "$#" -ge 2 ] || usage; DURATION_SEC="$2"; shift 2 ;;
        --devices) [ "$#" -ge 2 ] || usage; DEVICES="$2"; shift 2 ;;
        --exclude) [ "$#" -ge 2 ] || usage; EXCLUDE="$2"; shift 2 ;;
        --format) [ "$#" -ge 2 ] || usage; FORMAT="$2"; shift 2 ;;
        *) usage ;;
    esac
done

[ -x "${JOHN_BIN}" ] || { echo "john binary not found: ${JOHN_BIN}" >&2; exit 1; }

ALL_DEVICES=$(nvidia-smi --query-gpu=index --format=csv,noheader,nounits 2>/dev/null | sed 's/[[:space:]]//g' | awk 'NF' | paste -sd, -)
[ -n "${ALL_DEVICES}" ] || { echo "nvidia-smi found no NVIDIA GPUs" >&2; exit 1; }

DEVICES=$(normalize_list "${DEVICES}")
EXCLUDE=$(normalize_list "${EXCLUDE}")
SELECTED="${DEVICES}"
if [ -z "${SELECTED}" ]; then
    SELECTED="${ALL_DEVICES}"
fi

FINAL=""
for id in $(echo "${SELECTED}" | tr ',' ' '); do
    [ -n "${id}" ] || continue
    if contains_csv "${id}" "${EXCLUDE}"; then
        continue
    fi
    if [ -z "${FINAL}" ]; then
        FINAL="${id}"
    else
        FINAL="${FINAL},${id}"
    fi
done

[ -n "${FINAL}" ] || { echo "no NVIDIA GPUs selected after filters" >&2; exit 1; }

JOHN_DEVICES=""
for id in $(echo "${FINAL}" | tr ',' ' '); do
    opencl_id=$((id + 1))
    if [ -z "${JOHN_DEVICES}" ]; then
        JOHN_DEVICES="${opencl_id}"
    else
        JOHN_DEVICES="${JOHN_DEVICES},${opencl_id}"
    fi
done

echo "loader=john"
echo "selected_gpus=${FINAL}"
echo "john_devices=${JOHN_DEVICES}"

cd "${JOHN_DIR}"

ensure_opencl_ready || exit 1

choose_format() {
    if [ -n "${FORMAT}" ]; then
        echo "${FORMAT}"
        return 0
    fi
    for candidate in sha512crypt-opencl pbkdf2-hmac-sha512-opencl 7z-opencl sha256crypt-opencl md5crypt-opencl; do
        if ./john --test=1 --format="${candidate}" --devices="${JOHN_DEVICES}" >/dev/null 2>&1; then
            echo "${candidate}"
            return 0
        fi
    done
    return 1
}

CHOSEN_FORMAT=$(choose_format) || {
    echo "no suitable john OpenCL format found" >&2
    ./john --list=opencl-devices >&2 || true
    exit 1
}

run_john_loop() {
    opencl_id="$1"
    deadline="$2"
    round=0
    while :; do
        now=$(date +%s)
        remaining=$((deadline - now))
        if [ "${remaining}" -le 0 ]; then
            break
        fi
        round=$((round + 1))
        slice="${remaining}"
        if [ "${slice}" -gt "${TEST_SLICE_SECONDS}" ]; then
            slice="${TEST_SLICE_SECONDS}"
        fi
        echo "device=${opencl_id} round=${round} remaining_sec=${remaining} slice_sec=${slice}"
        ./john --test="${slice}" --format="${CHOSEN_FORMAT}" --devices="${opencl_id}" || return 1
    done
}

PIDS=""
cleanup() {
    rc=$?
    trap - EXIT INT TERM
    for pid in ${PIDS}; do
        kill "${pid}" 2>/dev/null || true
    done
    for pid in ${PIDS}; do
        wait "${pid}" 2>/dev/null || true
    done
    exit "${rc}"
}
trap cleanup EXIT INT TERM

echo "format=${CHOSEN_FORMAT}"
echo "target_seconds=${DURATION_SEC}"
echo "slice_seconds=${TEST_SLICE_SECONDS}"
DEADLINE=$(( $(date +%s) + DURATION_SEC ))
_first=1
for opencl_id in $(echo "${JOHN_DEVICES}" | tr ',' ' '); do
    [ "${_first}" = "1" ] || sleep 3
    _first=0
    run_john_loop "${opencl_id}" "${DEADLINE}" &
    pid=$!
    PIDS="${PIDS} ${pid}"
done
FAIL=0
for pid in ${PIDS}; do
    wait "${pid}" || FAIL=$((FAIL+1))
done
[ "${FAIL}" -eq 0 ] || { echo "john: ${FAIL} device(s) failed" >&2; exit 1; }