#!/bin/bash
#
# Copyright (C) 2025 MOXA Inc. All rights reserved.
# This software is distributed under the terms of the MOXA SOFTWARE NOTICE.
# See the file LICENSE for details.
#
# Authors:
#       2025  Elvis Yao  <elviscw.yao@moxa.com>

set -e

PKG_MAN=""
LOG_FILE=""
LOG_DIRNAME="$(hostname)_$(date +"%Y%m%d_%H%M%S")"
LOG_DIR="/var/log/mx_diagnostic/${LOG_DIRNAME}"
DIAGNOSTIC_LOG_PATH="${LOG_DIR}/diagnostic.log"
BENCHMARK_LOG_PATH="${LOG_DIR}/benchmark.log"
HWINFO_LOG_PATH="${LOG_DIR}/hwinfo.log"
SYSTEM_LOG_PATH="${LOG_DIR}/system.log"
HWINFO_LOG_PREFIX="[diagnostic][hwinfo]"
SYSTEM_LOG_PREFIX="[diagnostic][system]"
BENCHMARK_LOG_PREFIX="[diagnostic][benchmark]"

help_page() {
        cat <<EOF
Usage:  [option]
        mx-diagnostic.sh
Options:
        --help: show the help page
        --benchmark: execute stress test and benchmark test suite
        --hwinfo: extract detailed hardware information of the machine
        --system: collect all system information log

Without option to run the all diagnostic processes.
EOF
}

log_command() {
        local cmd_output
        local cur_timestamp

        # capture both stdout & stderr
        cmd_output="$("$@" 2>&1)"
        cur_timestamp="$(date +"%Y-%m-%d %H:%M:%S")"

        # shellcheck disable=SC2129
        {
                echo "[$cur_timestamp] Running: $*" >>"$LOG_FILE"
                echo "$cmd_output" >>"$LOG_FILE"
                echo "----------------------------------------" >>"$LOG_FILE"
        } >>"$LOG_FILE"
}

run_benchmark_prepare() {
        LOG_FILE="$BENCHMARK_LOG_PATH"
        touch "$LOG_FILE"

        echo "${BENCHMARK_LOG_PREFIX} Install necessary packages..."
        if [ "$PKG_MAN" == "apt" ]; then
                log_command apt-get install unzip libncurses5 sysbench fio -y
        elif [ "$PKG_MAN" == "yum" ]; then
                log_command yum install unzip libncurses5 sysbench fio -y
        else
                log_command echo "Package manager is not support."
                exit 1
        fi
}

run_benchmark() {
        echo "${BENCHMARK_LOG_PREFIX} Start to execute system benchmark..."
        run_benchmark_prepare

        # run sysbench
        echo "${BENCHMARK_LOG_PREFIX} Run sysbench for cpu/memory/disk benchmark..."
        log_command sysbench cpu --cpu-max-prime=20000 run
        log_command sysbench memory --memory-total-size=2G run
        log_command sysbench threads --threads=8 run
        log_command sysbench mutex --mutex-num=5000 run
        sysbench fileio --file-total-size=5G prepare >/dev/null
        log_command sysbench fileio --file-total-size=5G --file-test-mode=rndrw run
        sysbench fileio --file-total-size=5G cleanup >/dev/null

        # run fio
        echo "${BENCHMARK_LOG_PREFIX} Run fio for disk benchmark..."
        log_command fio --name=test --size=1G --rw=randrw --bs=4k --iodepth=64 --runtime=60 --numjobs=4 --group_reporting

        # run passmark
        echo "${BENCHMARK_LOG_PREFIX} Download and run passmark..."
        if [ ! -f "/opt/passmark/PerformanceTest/pt_linux_x64" ]; then
                wget https://www.passmark.com/downloads/pt_linux_x64.zip -O /tmp/pt_linux.zip >/dev/null
                unzip /tmp/pt_linux.zip -d /opt/passmark >/dev/null
        fi

        /opt/passmark/PerformanceTest/pt_linux_x64 -r 3 >/dev/null
        log_command cat results_all.yml

        echo "${BENCHMARK_LOG_PREFIX} Done"
}

run_system_prepare() {
        LOG_FILE="$SYSTEM_LOG_PATH"
        touch "$LOG_FILE"

        echo "${SYSTEM_LOG_PREFIX} Install necessary packages..."
        if [ "$PKG_MAN" == "apt" ]; then
                log_command apt-get install net-tools -y
        elif [ "$PKG_MAN" == "yum" ]; then
                log_command yum install net-tools -y
        else
                log_command echo "Package manager is not support."
                exit 1
        fi
}

run_system() {
        echo "${SYSTEM_LOG_PREFIX} Start to collect hardware information..."
        run_system_prepare

        # kernel message
        echo "${SYSTEM_LOG_PREFIX} Collect kernel message..."
        log_command dmesg
        log_command dmesg -l err

        # system info
        echo "${SYSTEM_LOG_PREFIX} Collect system information..."
        log_command hostname
        log_command uname -a

        # systemd information
        echo "${SYSTEM_LOG_PREFIX} Collect running services information..."
        log_command systemctl list-units --type=service --no-pager
        log_command journalctl -n 100 --no-pager

        # network information
        echo "${SYSTEM_LOG_PREFIX} Collect network information..."
        log_command ip a
        log_command ip route
        log_command netstat -tulnp

        # network information
        echo "${SYSTEM_LOG_PREFIX} Collect DMI table information..."
        log_command dmidecode

        echo "${SYSTEM_LOG_PREFIX} Done"
}

run_hwinfo_prepare() {
        LOG_FILE="$HWINFO_LOG_PATH"
        touch "$LOG_FILE"

        echo "${HWINFO_LOG_PREFIX} Install necessary packages..."
        if [ "$PKG_MAN" == "apt" ]; then
                log_command apt-get install lshw util-linux usbutils pciutils smartmontools -y
        elif [ "$PKG_MAN" == "yum" ]; then
                log_command yum install lshw util-linux usbutils pciutils smartmontools -y
        else
                log_command echo "Package manager is not support."
                exit 1
        fi
}

run_hwinfo() {
        echo "${HWINFO_LOG_PREFIX} Start to collect hardware information..."
        run_hwinfo_prepare

        # all hw info
        echo "${HWINFO_LOG_PREFIX} Collect all hardware information by lshw..."
        log_command lshw

        # cpu info
        echo "${HWINFO_LOG_PREFIX} Collect cpu information..."
        log_command lscpu

        # memory info
        echo "${HWINFO_LOG_PREFIX} Collect memory information..."
        log_command free -h
        log_command cat /proc/meminfo

        # disk device info
        echo "${HWINFO_LOG_PREFIX} Collect disk information..."
        log_command lsblk
        log_command df -h
        blk_devices=$(lsblk -dno NAME,TYPE,TRAN | grep -E 'disk' | grep -Ev 'usb' | awk '{print "/dev/"$1}')
        # Loop through devices and scan each
        for DEV in $blk_devices; do
                log_command smartctl -a "$DEV"
        done

        # pci device info
        echo "${HWINFO_LOG_PREFIX} Collect PCI devices information..."
        log_command lspci -tv
        log_command lspci -vvv

        # usb device info
        echo "[diagnostic][hwinfo] Collect USB devices information..."
        log_command lsusb -tvvv

        echo "${HWINFO_LOG_PREFIX} Done"
}

run_all() {
        run_hwinfo
        run_system
        run_benchmark
}

run_prepare() {
        mkdir -p "$LOG_DIR"
        LOG_FILE="$DIAGNOSTIC_LOG_PATH"
        touch "$LOG_FILE"

        if command -v apt-get &>/dev/null; then
                PKG_MAN="apt"
                log_command apt-get update
                if command -v lsb_release -a &>/dev/null; then
                        log_command apt-get install lsb-release libcurl4 -y
                fi
        elif command -v yum &>/dev/null; then
                PKG_MAN="yum"
                log_command yum update
                if command -v lsb_release -a &>/dev/null; then
                        log_command yum install lsb-release libcurl4 -y
                fi
        else
                log_command echo "Package manager is not support."
                exit 1
        fi

        # get os info
        log_command lsb_release -a
}

main() {
        local param

        if [ $# -eq 0 ]; then
                echo -e "No options provided.\nTo run all diagnostic processes: hwinfo/system/benchmark."
                run_prepare
                run_all
        else
                param="$1"
                case "$param" in
                --help)
                        help_page
                        exit 0
                        ;;
                --benchmark)
                        run_prepare
                        run_benchmark
                        ;;
                --system)
                        run_prepare
                        run_system
                        ;;
                --hwinfo)
                        run_prepare
                        run_hwinfo
                        ;;
                *)
                        echo "Invalid option was provided: $param"
                        help_page
                        exit 22
                        ;;
                esac
        fi

        echo "[diagnostic] Diagnostic logs collected in: $LOG_DIR"
}

main "$@"
