#!/bin/sh

# ==============================================================================
# NOR Flash Lifetime Stress Test for Embedded Linux (BusyBox)
#
# Version: 1.2
#
# Description:
# Performs a continuous write/read/verify cycle on a NOR flash MTD partition.
# Creates a status file in /tmp for monitoring during long-running tests.
#
# WARNING: THIS SCRIPT IS DESTRUCTIVE. IT WILL ERASE ALL DATA ON THE
#          SPECIFIED FLASH DEVICE. USE WITH EXTREME CAUTION.
#
# How to monitor:
# In another shell, run: cat /tmp/nor_stress_test_status.log
# ==============================================================================

# --- CONFIGURATION ---
FLASH_DEVICE="/dev/mtdblock5"
TEST_FILE_SIZE_KB=64
TEMP_DIR="/tmp"
# --- END OF CONFIGURATION ---


# --- SCRIPT LOGIC ---
SOURCE_FILE="${TEMP_DIR}/stress_source.dat"
READBACK_FILE="${TEMP_DIR}/stress_readback.dat"
SOURCE_MD5_FILE="${TEMP_DIR}/source.md5"
READBACK_MD5_FILE="${TEMP_DIR}/readback.md5"
STATUS_FILE="${TEMP_DIR}/nor_stress_test_status.log"

START_TIME=$(date)
passed_cycles=0
cycle_count=1

# Function to write the current status to a log file
update_status() {
  current_status_msg="$1"
  # Overwrite the status file with the latest info
  {
    printf "--- NOR Flash Stress Test Status ---\n"
    printf "Test Start Time:     %s\n" "${START_TIME}"
    printf "Last Update:         %s\n" "$(date)"
    printf "Total Cycles Passed: %d\n" "${passed_cycles}"
    printf "Current Cycle:       %d\n" "${cycle_count}"
    printf "Current Status:      %s\n" "${current_status_msg}"
  } > "${STATUS_FILE}"
}

# Function to clean up temporary files on exit
cleanup() {
  printf "\nCleaning up temporary files...\n"
  rm -f "${SOURCE_FILE}" "${READBACK_FILE}" "${SOURCE_MD5_FILE}" "${READBACK_MD5_FILE}"
  printf "Cleanup complete. Exiting.\n"
}

trap cleanup INT TERM EXIT

# --- Initial Sanity Checks ---
printf "============================================================\n"
printf "       NOR Flash Stress Test Script (v1.2)\n"
printf "============================================================\n"
printf "Target Device:     %s\n" "${FLASH_DEVICE}"
printf "Test Data Size:    %d KB\n" "${TEST_FILE_SIZE_KB}"
printf "Status Log File:   %s\n" "${STATUS_FILE}"
printf "------------------------------------------------------------\n"
printf "!!! WARNING: THIS WILL DESTROY ALL DATA ON %s !!!\n" "${FLASH_DEVICE}"
printf "------------------------------------------------------------\n"
printf "Type 'YES' to continue: "
read -r confirmation
if [ "${confirmation}" != "YES" ]; then
  printf "Confirmation not received. Aborting.\n"
  exit 1
fi

if [ ! -b "${FLASH_DEVICE}" ]; then
  printf "Warning: Device %s not found or is not a character device.\n" "${FLASH_DEVICE}"
fi

printf "Starting stress test. Press Ctrl+C to stop.\n"
printf "You can monitor progress by running 'cat %s'\n" "${STATUS_FILE}"

# --- Main Test Loop ---
while true; do
  printf "\n----- Cycle %d -----\n" "${cycle_count}"

  # 1. Create source file
  update_status "Creating random source file"
  printf "[1/5] Creating %d KB random source file...\n" "${TEST_FILE_SIZE_KB}"
  dd if=/dev/urandom of="${SOURCE_FILE}" bs=1k count="${TEST_FILE_SIZE_KB}" 2>/dev/null
  if [ $? -ne 0 ]; then
    update_status "ERROR: Failed to create source file. Test stopped."
    printf "Error: Failed to create source file. Aborting.\n"
    exit 1
  fi

  # 2. Calculate source MD5
  update_status "Calculating source MD5 hash"
  printf "[2/5] Calculating source file MD5 hash...\n"
  md5sum "${SOURCE_FILE}" > "${SOURCE_MD5_FILE}"
  read -r source_md5 _ < "${SOURCE_MD5_FILE}"
  printf "      Source MD5: %s\n" "${source_md5}"

  # 3. Write to flash
  update_status "Writing data to ${FLASH_DEVICE}"
  printf "[3/5] Writing data to %s...\n" "${FLASH_DEVICE}"
  dd if="${SOURCE_FILE}" of="${FLASH_DEVICE}" bs=1k 2>/dev/null
  if [ $? -ne 0 ]; then
    update_status "ERROR: Failed to write to flash. Test stopped."
    printf "Error: Failed to write to %s. Aborting.\n" "${FLASH_DEVICE}"
    exit 1
  fi

  # 4. Read back from flash
  update_status "Reading data back from ${FLASH_DEVICE}"
  printf "[4/5] Reading data back from %s...\n" "${FLASH_DEVICE}"
  dd if="${FLASH_DEVICE}" of="${READBACK_FILE}" bs=1k count="${TEST_FILE_SIZE_KB}" 2>/dev/null
  if [ $? -ne 0 ]; then
    update_status "ERROR: Failed to read from flash. Test stopped."
    printf "Error: Failed to read from %s. Aborting.\n" "${FLASH_DEVICE}"
    exit 1
  fi

  # 5. Verify data
  update_status "Verifying data integrity"
  printf "[5/5] Verifying data integrity...\n"
  md5sum "${READBACK_FILE}" > "${READBACK_MD5_FILE}"
  read -r readback_md5 _ < "${READBACK_MD5_FILE}"
  printf "      Readback MD5: %s\n" "${readback_md5}"

  if [ "${source_md5}" = "${readback_md5}" ]; then
    passed_cycles=$((passed_cycles + 1))
    printf "      SUCCESS: Hashes match. Cycle %d passed.\n" "${cycle_count}"
    update_status "Cycle passed. Waiting for next cycle."
  else
    failure_msg="FAILURE: HASH MISMATCH on cycle ${cycle_count}"
    update_status "${failure_msg}"
    printf "      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n"
    printf "      !!! %s !!!\n" "${failure_msg}"
    printf "      !!! Source MD5:   %s\n" "${source_md5}"
    printf "      !!! Readback MD5: %s\n" "${readback_md5}"
    printf "      !!! Flash may be failing. Aborting test. !!!\n"
    printf "      !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n"
    mv "${SOURCE_FILE}" "${TEMP_DIR}/FAILED_source_cycle_${cycle_count}.dat"
    mv "${READBACK_FILE}" "${TEMP_DIR}/FAILED_readback_cycle_${cycle_count}.dat"
    exit 1
  fi

  cycle_count=$((cycle_count + 1))
done