#! /bin/bash
# hpc-testing
# Copyright (C) 2025 SUSE LLC
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

nvme_client_pre_cleanup()
{
    mount | grep /dev/nvme | awk '{ print $1}' | xargs -I - umount - 2>/dev/null
    nvme disconnect -n testnq || true
    return 0
}

nvme_setup_server()
{
    local server_ip=$1
    local LOOPD

    # Pre cleanup. Ignore errors here
    set +e
    umount /tmp/hpc-test.mount 2>/dev/null
    rm -Rf /tmp/hpc-test.mount /tmp/hpc-test.io 2>/dev/null
    losetup -a | grep hpc-test.io | sed -e s/:.*// | xargs losetup -d
    modprobe nvmet
    nvmetcli clear
    set -e

    modprobe nvmet_rdma
    LOOPD=$(losetup -f)
    dd if=/dev/zero of=/tmp/hpc-test.io bs=1M count=256
    losetup ${LOOPD} /tmp/hpc-test.io
    mkfs.ext3 ${LOOPD}
    mkdir /tmp/hpc-test.mount/
    mount ${LOOPD} /tmp/hpc-test.mount/
    dd if=/dev/urandom bs=1M count=64 of=/tmp/hpc-test.mount/input
    umount  /tmp/hpc-test.mount
    sed -i -e s/@MYIP@/"$server_ip"/ -e s%@BLK@%${LOOPD}% hpc-nvmet.json
    nvmetcli restore hpc-nvmet.json
}

nvme_run_client()
{
    local server_ip=$1
    local block_device
    local nvme_dev_pre nvme_dev_post
    set +e
    umount /tmp/srp-hpc-test 2>/dev/null
    rm -Rf /tmp/srp-hpc-test 2>/dev/null
    set -e

    modprobe nvme_rdma
    mkdir /tmp/srp-hpc-test
    rm -f /etc/nvme/hostid
    nvme discover -t rdma -a "$server_ip" -s 4420
    nvme_dev_pre=$(lsblk | grep nvme | cut -d' ' -f1 | sort)
    nvme connect -t rdma -n testnqn -a "$server_ip" -s 4420
    nvme_dev_post=$(lsblk | grep nvme | cut -d' ' -f1 | sort)
    sleep 1
    block_device=$(comm -13 <(echo $nvme_dev_pre) <(echo $nvme_dev_post))
    echo $block_device
    if [ "$block_device" == "" ]; then
        echo "No block device found"
        exit 1
    fi
    mount /dev/$block_device /tmp/srp-hpc-test
    cp -R /tmp/srp-hpc-test/input /tmp/srp-hpc-test/output
    diff -q /tmp/srp-hpc-test/input /tmp/srp-hpc-test/output
    umount /tmp/srp-hpc-test
    nvme disconnect -n testnqn
    ! (lsblk | grep nvme)
}

nvme_test_and_clear_server()
{
    nvmetcli clear
    mount -o loop /tmp/hpc-test.io /tmp/hpc-test.mount
    diff -q /tmp/hpc-test.mount/input /tmp/hpc-test.mount/output
    umount /tmp/hpc-test.mount
    losetup -a | grep hpc-test.io | sed -e s/:.*// | xargs losetup -d
}

test_nvme(){
    local server=$1
    local server_ip=$2
    local client=$3

    tpq $server 'cat > hpc-nvmet.json' < helpers/ib/nvmet.json

    tp_fun $client nvme_client_pre_cleanup
    tp_fun $server nvme_setup_server $server_ip
    tp_fun $client nvme_run_client $server_ip
    tp_fun $server nvme_test_and_clear_server
}
