#!/bin/bash
Z_WORKDIR="$(realpath "$(dirname "${0}")")"
source "${Z_WORKDIR}/ansihelper.sh"

drives=($(ls -1 /dev/sd?))

function zcmd() {
                if [ -n "${DEBUG}" ]; then
                        echo "-- Executing ${*} ..." >&2
                fi
                "${@}"
}

function zsmart() {
                local drive sopts;
                sopts=("${@}")
                drive="${sopts[$((${#sopts[@]}-1))]}"
                unset -v "sopts[$((${#sopts[@]}-1))]"
                case "${drive}" in
                        "/dev/sdi")
                                        sopts+=("-dsat")
                                        ;;
                esac
                zcmd sudo smartctl "${sopts[@]}" "${drive}";
}

function zgetsts() {
        local sts;
        sts="$(sudo smartctl -c "${1}" | grep "Self-test execution status" | cut -d'(' -f2 | cut -d')' -f1 | sed 's| ||g')"
        case "${sts}" in
                0)
                        echo "No Self-test in Progress"
                        ;;
                24?)
                        echo "Extended Self-test in Progress (${sts:2:1}0% remaining)"
                        ;;
                *)
                        echo "Unknown Self-test status code ${sts}"
                        ;;
        esac
}

function zgetpoh() {
        local poh;
        poh=$(zsmart -A "${1}" | grep Power_On | rev | cut -d'-' -f1 | rev | sed 's| ||g' | cut -d'(' -f1)
        if [ ${poh} -gt 65535 ]; then
                poh="${poh} (SMART Self-test Log POH: $((poh - 65535)))"
        fi
        echo "${poh}"
}

function smartctl_all_drives_cond() {
                if [ "${2}"  == "status" ]; then
                        if [ -e "${1}" ]; then
                                d="${1}";
                                str_color "green" "-- Working on ${d} ..." >&2
                                echo ""
                                zsmart -H "${d}"
                                str_color "yellow" "-- Power On Hours: $(str_color "green" "$(zgetpoh "${d}")")"
                                echo ""
                                str_color "yellow" "-- Self-test status: $(str_color "green" "$(zgetsts "${d}")")"
                                echo ""
                                zsmart --log=xselftest "${d}"| grep -A3 "SMART Extended Self-test Log"
                        else
                                for d in "${drives[@]}"; do
                                        str_color "green" "-- Working on ${d} ..." >&2
                                        echo ""
                                        zsmart -H "${d}"
                                        str_color "yellow" "-- Power On Hours: $(str_color "green" "$(zgetpoh "${d}")")"
                                        echo ""
                                        str_color "yellow" "-- Self-test status: $(str_color "green" "$(zgetsts "${d}")")"
                                        echo ""
                                        zsmart --log=xselftest "${d}"| grep -A3 "SMART Extended Self-test Log"
                                done
                        fi
                else
                        if [ -e "${1}" ]; then
                                d="${1}";
                                shift;
                                str_color "green" "-- Working on ${d} ..." >&2
                                echo ""
                                zsmart "${@}" "${d}";
                        else
                                shift;
                                for d in "${drives[@]}"; do
                                        str_color "green" "-- Working on ${d} ..." >&2
                                        echo ""
                                        zsmart "${@}" "${d}";
                                done
                        fi
                fi
}

case "${1}" in
        "extended")
                smartctl_all_drives_cond "${2}" --test=long
                ;;
        "short")
                smartctl_all_drives_cond "${2}" --test=short
                ;;
        "rawcmd")
                shift;
                smartctl_all_drives_cond "${@}"
                ;;

        "cancel")
                smartctl_all_drives_cond "${2}" -X
                ;;
        "status")
                smartctl_all_drives_cond "${2}" "${1}"
                ;;
        *)
                echo "Usage: ${0} [status|extended|short|cancel] (/dev/drive)"
                exit 1;
                ;;
esac