# Copyright 2024 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

apiVersion: v1
kind: ConfigMap
metadata:
  name: file-generator-config
data:
  generate_files.sh: |
    #!/bin/bash

    RETRY_COUNTER=0
    MAX_RETRIES=5
    SLEEP_TIME=1

    # Check if OUTPUT_DIR environment variable is set, otherwise use default
    if [ -z "${OUTPUT_DIR}" ]; then
        OUTPUT_DIR="/data"  # Default to /data inside the container
        echo "OUTPUT_DIR not set, using default: ${OUTPUT_DIR}"
    fi

    echo "Loading Data in to $OUTPUT_DIR"
    echo "Waiting for Parallelstore mount..."

    OUTPUT_BASE_DIR=$(dirname "${OUTPUT_DIR}")

    while [ $RETRY_COUNTER -lt $MAX_RETRIES ]; do
      if mountpoint -q $OUTPUT_BASE_DIR; then
        echo "[$(date +%Y-%m-%d_%H:%M:%S)] Parallelstore mount detected at $OUTPUT_BASE_DIR"
        break
      fi
      echo "[$(date +%Y-%m-%d_%H:%M:%S)] Parallelstore mount not detected, retrying in $SLEEP_TIME seconds..."
      sleep $SLEEP_TIME
      COUNTER=$((COUNTER+1))
      SLEEP_TIME=$((SLEEP_TIME * 2)) # Exponential backoff
      if [ $RETRY_COUNTER -eq $MAX_RETRIES ]; then
        echo "[$(date +%Y-%m-%d_%H:%M:%S)] Failed to detect Parallelstore mount at $OUTPUT_BASE_DIR after multiple retries"
        # Call cleanup function if needed
        cleanup  # Make sure the 'cleanup' function is defined if you use it.
        exit 1
      fi
    done


    # Create or clean directory for test files (using /data)
    if [ -d "${OUTPUT_DIR}" ]; then
        echo "Cleaning existing directory: ${OUTPUT_DIR}"
        rm -rf "${OUTPUT_DIR}"/*
    else
        echo "Creating directory: ${OUTPUT_DIR}"
    fi
    mkdir -p "${OUTPUT_DIR}"

    # Function to generate a single file using FIO
    generate_file() {
        file_num=$1
        size=$2
        size_arg="${size_kb}K"
        fio --name=generate_file_${file_num} \
            --ioengine=libaio \
            --rw=write \
            --bs=4k \
            --direct=1 \
            --size=${size_arg} \
            --filename="${OUTPUT_DIR}/testfile_${file_num}" \
            --thread \
            --group_reporting \
            --minimal
            # >/dev/null 2>&1  # Redirect output to /dev/null
    }
    start_time=$(date +%s)
    echo "Starting file generation..."

    # Generate 1000 files with random sizes between 10KB and 1MB
    for file_num in $(seq 1 1000); do
        # Generate random size between 10KB and 1MB (in KB)
        size_kb=$(( (RANDOM % 990) + 10 )) # Generates 20-1000KB

        # Convert to MB for fio (with two decimal places)
        size=$(awk "BEGIN {printf \"%.2f\", $size_kb / 1024}")

        # Show progress every 100 files
        if [ $((file_num % 100)) -eq 0 ]; then
            echo "Generated $file_num files..."
        fi

        generate_file "$file_num" "$size" &
    done

    wait

    end_time=$(date +%s)
    elapsed_time=$((end_time - start_time))
    echo "File generation complete!"
    echo "Total time taken: $elapsed_time seconds"

    # Display summary of generated files (adapted for /data)
    total_size=$(du -sh "${OUTPUT_DIR}" | cut -f1)
    echo "Total size of generated files: $total_size"
    echo "File size distribution:"
    ls -l "${OUTPUT_DIR}" | awk '{print $5}' | awk '{ sum += $1; n++ } END { print "Average file size: " sum/n/1024/1024 " MB" }'

    # Optional: Show detailed distribution
    echo -e "\nDetailed size distribution:"
    ls -l "${OUTPUT_DIR}" | awk '{print int($5/1024/1024)"MB"}' | sort | uniq -c
