#!/bin/bash

set -e

# Valeurs par défaut
FILESIZE="1G"
FILECOUNT=1000
FILEMODE="fixed"
LINES=1
FLUSH_CACHE="no"
RANDOM_WRITE="no"

# Fonction pour convertir "M:SS.xxx" en secondes float
convert_time_to_sec() {
  local raw="$1"
  IFS=":" read -r min sec <<< "$raw"
  echo | awk -v m="${min:-0}" -v s="${sec:-0}" 'BEGIN {
    total = (m * 60) + s;
    if (total <= 0) {
      print "0.001";
    } else {
      print total;
    }
  }'
}

function show_help {
    echo "Usage: $0 [OPTIONS]"
    echo ""
    echo "Options disponibles :"
    echo "  --size=SIZE         Taille du fichier dd (ex: 1G, 500M). Défaut: 1G"
    echo "  --count=N           Nombre de petits fichiers. Défaut: 1000"
    echo "  --mode=MODE         Mode de génération : fixed ou random. Défaut: fixed"
    echo "  --lines=N           Nombre de lignes pour les fichiers en mode fixed. Défaut: 1"
    echo "  --flush-cache       Vide les caches Linux entre écriture et lecture (nécessite root)"
    echo "  --random-write      Utilise /dev/urandom pour écrire au lieu de /dev/zero"
    echo "  --help              Affiche ce message d'aide"
    echo ""
    echo "Exemples:"
    echo "# Benchmark basique"
    echo "./benchmark_disque.sh"
    echo ""
    echo "# Test avec un fichier de 2G et 5000 petits fichiers aléatoires"
    echo "./benchmark_disque.sh --size=2G --count=5000 --mode=random"
    echo ""
    echo "# Test avec fichiers de 10 lignes"
    echo "./benchmark_disque.sh --mode=fixed --lines=10"
    exit 0
}

# Parsing des arguments
for arg in "$@"; do
    case $arg in
        --size=*)
            FILESIZE="${arg#*=}"
            ;;
        --count=*)
            FILECOUNT="${arg#*=}"
            ;;
        --mode=*)
            FILEMODE="${arg#*=}"
            ;;
        --lines=*)
            LINES="${arg#*=}"
            ;;
        --flush-cache)
            FLUSH_CACHE="yes"
            ;;
        --random-write)
            RANDOM_WRITE="yes"
            ;;
        --help)
            show_help
            ;;
        *)
            echo "Argument inconnu : $arg"
            show_help
            ;;
    esac
done

# Définir la source d'écriture
if [[ "$RANDOM_WRITE" == "yes" ]]; then
    WRITE_SRC="/dev/urandom"
else
    WRITE_SRC="/dev/zero"
fi

# Répertoires
BENCHDIR="./testbench"
TESTFILE="$BENCHDIR/dd_testfile"
SMALLDIR="$BENCHDIR/smallfiles"

mkdir -p "$SMALLDIR"

echo "=== Test d'écriture séquentielle avec dd ($FILESIZE, source: $(basename $WRITE_SRC), sync forcée) ==="
WRITE_TIME_SEC=$( (time dd if="$WRITE_SRC" of="$TESTFILE" bs="$FILESIZE" count=1 conv=fdatasync status=none) 2>&1 | grep real | awk '{print $2}' | sed 's/m/:/;s/s//')
WRITE_SEC=$(convert_time_to_sec "$WRITE_TIME_SEC")
echo -n "Durée : $WRITE_TIME_SEC"
echo " (~$(awk -v size="$FILESIZE" -v sec="$WRITE_SEC" 'BEGIN {
  gsub(/[A-Za-z]/, "", size)
  if (index(tolower(ARGV[1]), "g")) mult=1024; else mult=1;
  printf "%.2f", (size * mult) / sec
}') MB/s)"

if [[ "$FLUSH_CACHE" == "yes" ]]; then
    echo -e "\nVidage du cache Linux (nécessite root)..."
    sync && echo 3 > /proc/sys/vm/drop_caches
fi

echo -e "\n=== Test de lecture séquentielle ==="
READ_TIME_SEC=$( (time dd if="$TESTFILE" of=/dev/null bs="$FILESIZE" status=none) 2>&1 | grep real | awk '{print $2}' | sed 's/m/:/;s/s//')
READ_SEC=$(convert_time_to_sec "$READ_TIME_SEC")
echo -n "Durée : $READ_TIME_SEC"
echo " (~$(awk -v size="$FILESIZE" -v sec="$READ_SEC" 'BEGIN {
  gsub(/[A-Za-z]/, "", size)
  if (index(tolower(ARGV[1]), "g")) mult=1024; else mult=1;
  printf "%.2f", (size * mult) / sec
}') MB/s)"

if (( $(echo "$READ_SEC < 0.01" | bc -l) )); then
    echo "(⚠️ lecture probablement en cache — résultat peu fiable)"
fi

echo -e "\n=== Test IOPS sur $FILECOUNT fichiers ($FILEMODE) ==="

echo -n "Création : "
CREATE_TIME_SEC=$( (time for i in $(seq 1 "$FILECOUNT"); do
    if [[ "$FILEMODE" == "random" ]]; then
        head -c $((RANDOM % 4096 + 128)) /dev/urandom > "$SMALLDIR/file_$i"
    else
        yes "Ligne test $i" | head -n "$LINES" > "$SMALLDIR/file_$i"
    fi
done) 2>&1 | grep real | awk '{print $2}' | sed 's/m/:/;s/s//')
CREATE_SEC=$(convert_time_to_sec "$CREATE_TIME_SEC")
CREATE_IOPS=$(awk -v total="$FILECOUNT" -v sec="$CREATE_SEC" 'BEGIN { printf "%.1f", total / sec }')
echo "$CREATE_TIME_SEC (~$CREATE_IOPS IOPS)"

echo -n "Lecture : "
READ_SMALL_SEC=$( (time for i in $(seq 1 "$FILECOUNT"); do
    cat "$SMALLDIR/file_$i" > /dev/null
done) 2>&1 | grep real | awk '{print $2}' | sed 's/m/:/;s/s//')
READ_SMALL_SEC_VAL=$(convert_time_to_sec "$READ_SMALL_SEC")
READ_SMALL_IOPS=$(awk -v total="$FILECOUNT" -v sec="$READ_SMALL_SEC_VAL" 'BEGIN { printf "%.1f", total / sec }')
echo "$READ_SMALL_SEC (~$READ_SMALL_IOPS IOPS)"

echo -e "\nNettoyage..."
rm -rf "$BENCHDIR"

echo -e "\nBenchmark terminé."

