#!/bin/bash

# Created by jiaqiang.ye@lnxall.com
# Created by Artificial Intelligence
# 2025/02/24

# Function to check if the root filesystem is full
check_root_fs_full() {
    local threshold=95  # Consider root filesystem full if usage exceeds 92%
    local usage=$(df -h | grep '/dev/nvme0n1' | awk '{print $5}' | sed 's/%//g')

    if [ "$usage" -gt "$threshold" ]; then
        echo "Root filesystem is full! Usage: ${usage}%"
        return 1  # Return non-zero to indicate the filesystem is full
    fi
    return 0  # Return zero to indicate no issue
}

# Function to write a 32MB random file and check its MD5SUM
write_random_file() {
    local f="random_data_$1.bin"
    
    # Write 32MB of random data to the file and compute MD5SUM during writing
    local orig_md5=$(dd if=/dev/urandom bs=8192 count=4096 2>/dev/null | tee "$f" | md5sum - | awk '{print $1}')
    
    # Flush filesystem cache to ensure data is written to the storage device
    sync
	echo 3 > /proc/sys/vm/drop_caches
    
    # Check the size of the written file
    local size=$(stat --format="%s" "$f")
    
    if [ "$size" -ne $((32 * 1024 * 1024)) ]; then
        return 1  # Return non-zero to indicate a size mismatch
    fi
    
    # Get the MD5SUM of the written file
    local written_md5=$(md5sum "$f" | awk '{print $1}')
    
    # Compare the two MD5SUMs
    if [ "$orig_md5" == "$written_md5" ]; then
        return 0  # Return zero to indicate success
    fi
    return 2  # Return 2 to indicate a corrupted file
}

check_health() {
	local msg=$(sudo nvme smart-log /dev/nvme0n1 | grep -e 'percentage_used' -e 'critical_warning')

	echo ""	
	if [ $? -eq 0 ] ; then
		echo "$msg"
	else
		echo "failed to check health."
	fi
}

# emit current date-time in a function, so we can easily modify its format.
current_DT() {
    echo "$(date '+%Y-%m-%d %H:%M:%S')"
}

# Main loop: Stress test the Flash-Storage device
declare -i epoch=1 # We are using bash, this is valid
prev_files=0  # Track the number of files created in the previous epoch

systemctl stop emsd         2>/dev/null
systemctl stop emud         2>/dev/null
systemctl stop linit        2>/dev/null
systemctl stop lnxhostapd   2>/dev/null
mountpoint /mnt/ssd
if [ $? -ne 0 ] ; then
	echo "Error, SSD not found"
fi
mkdir -p /mnt/ssd/flash_stress_test
cd /mnt/ssd/flash_stress_test

while true; do
    # Print the start of the epoch with a date-time prefix
    echo "$(current_DT): Starting epoch $epoch (Files in previous epoch: $prev_files)..."
    
    declare -i file_count=1
    while true; do
        # Generate a dynamic status line that updates itself
        echo -n -e "\r$(current_DT): Generating random data, Epoch ${epoch}, file random_data_${file_count}.bin (${prev_files} total)..."

        # Write a random file and check its integrity
        write_random_file "$file_count"
        result=$?
        
        if [ $result -eq 1 ]; then
            # File size mismatch, check if the root filesystem is full
            check_root_fs_full
            if [ $? -ne 0 ]; then
                echo ""
                echo "$(current_DT): No space left on root filesystem. Deleting all files and starting a new epoch..."
                rm -f random_data_*.bin
                prev_files=$((file_count - 1))  # Store the number of files created in this epoch
                break
            fi
            # Root filesystem is not full, but file size mismatch indicates Flash-Storage failure
            echo ""
            echo "$(current_DT): Warning: File size mismatch detected, but root filesystem is not full!"
            echo "$(current_DT): Something is wrong with the Flash-Storage device. Terminating script."
            exit 1
        elif [ $result -eq 2 ]; then
            # MD5SUM mismatch, but there is still free space (Flash-Storage failure)
            echo ""
            echo "$(current_DT): Corrupted file detected! Flash-Storage device cannot correctly store random data."
            echo "$(current_DT): Job done. Terminating script."
            exit 0
        fi
        
        # Increment file count
        file_count+=1
    done
	check_health

	echo ""
    # Increment epoch number
    epoch+=1 # Valid for variables in bash script, declared with `declare -i`
done