#!/usr/bin/env bash
# This script is used to split large files into smaller chunks.
# It can be used to split files stored in s3 or locally.
#
# When splitting files from s3, it downloads the file, splits it into chunks,
# compress the chunks and uploads them to s3.
# It doesn't requires drive storage for downloaded S3 file locally, it streams it to the split command.
# The file is split on per-line basis.
#
# required commands: aws, zcat, split, gzip, pigz (optionally)
#

function usage {
    cat <<EOF
usage:
split_file_s3.sh -s <s3://bucket/source_file/path/filename.gz> -t <s3://bucket/destination/> [-z <uncompressed chunk size, default 50m>] [-c <compression level, default 5>] [-e <skip_header, default 0>] [-a <archiver, default gzip>]

Arguemnts:
  -s  source_file key, like s3://bucket/source_file/path/filename.gz or /path/to/local/file.gz
  -t  target path, like s3://bucket/destination/ or /path/to/local/splits/
  -z  uncompressed chunk size, default '50m' - 50 Megabytes
  -c  compression level (1-9), default 5
  -e  skip lines from file header, default 0
  -a  archiver (gzip, pigz), default gzip
  -v  verbose mode

Example:
  # Example 1: download from s3, split and upload to s3 (fastest compression using pigz, skip 1 line from header)
  ./split_file_s3.sh -s s3://dev-cucumbers/YouTubeAssetConflict/archives/2024-02-20/content_owner_asset_conflict_a2.ORCH.csv.gz -t s3://dev-cucumbers/YouTubeAssetConflict/splits/2024-02-20/ -a pigz -c 1
  Note: the tested processing speed of s3-to-s3 case was up to 18 Mb/sec

  # Example 2: split local file and save to local folder (medium compression using gzip, no header skip)
  ./split_file_s3.sh -s ./content_owner_asset_conflict_a2.ORCH.csv.gz -t ./files

  # Example 3: split local file and upload to s3 (max compression using gzip, no header skip, chnuk size 50 megabytes)
  ./split_file_s3.sh -s ./content_owner_asset_conflict_a2.ORCH.csv.gz -t s3://dev-cucumbers/YouTubeAssetConflict/splits/2024-02-20/ -z 50m -c 9
EOF
    exit 3
}

while getopts ":s:t:z:c:d:a:vh" opt; do
  case ${opt} in
    s )
      source_file=$OPTARG
      ;;
    t )
      target_path=$OPTARG
      ;;
    z )
      uncompressed_chunk_size=$OPTARG
      ;;
    c )
      compression_level=$OPTARG
      ;;
    a )
      archiver=$OPTARG
      ;;
    d )
      skip_header=$OPTARG
      ;;
    v )
      verbose='--verbose'
      ;;
    h )
      usage
      ;;
    \? )
      echo "Wrong arg: -$OPTARG" 1>&2
      echo "Type -h for help" 1>&2
      exit 4
      ;;
    : )
      echo "Arg -$OPTARG requires value" 1>&2
      echo "Type -h for help" 1>&2
      exit 4
      ;;
  esac
done

set -e
set -o pipefail

if [ -z "$source_file" ] ; then
      echo "Option -s is required and should contain full s3 path of the source_file file" 1>&2
      echo "Type -h for help" 1>&2
      exit 4
fi
if [ -z "$target_path" ] ; then
      echo "Option -t is required and should represent s3 destination path" 1>&2
      echo "Type -h for help" 1>&2
      exit 4
fi
if [[ $target_path != *\/ ]] ; then
      echo "Option -t is should ends with slash /" 1>&2
      exit 4
fi

verbose=${verbose:-}
skip_header=${skip_header:-0}
uncompressed_chunk_size=${uncompressed_chunk_size:-15m}
compression_level=${compression_level:-5}
archiver=${archiver:-gzip}
filename=`basename $source_file`

echo "Source: $source_file"
echo "Skip header lines: $skip_header"
echo "Filename: $filename"
echo "Target: $target_path"
echo "Uncompressed chunk size: $uncompressed_chunk_size"
echo "Archiver: $archiver"
echo "Compression level: $compression_level"

if [[ $source_file == s3:\/\/* ]]; then
  reader="aws s3 cp ${source_file} -"
else
  reader="cat $source_file"
fi

if [ $skip_header -ge 1 ]; then
  # skip header lines
  prepocessor="tail --lines=+$((skip_header+1))"
else
  prepocessor="cat"
fi

# check if file ends with .gz
if [[ $filename == *.gz ]]; then
  decompressor="zcat -"
else
  decompressor="cat"
fi

if [[ $archiver == "pigz" ]]; then
  compressor="pigz - -${compression_level}"
else
  compressor="gzip -c -${compression_level}"
fi

if [[ $target_path == s3:\/\/* ]]; then
  writer="aws s3 cp - ${target_path}\$FILE.gz"
else
  mkdir -p $target_path
  writer="cat > ${target_path}\$FILE.gz"
fi


$reader | \
  $decompressor | \
  $prepocessor | \
  split $verbose --numeric-suffixes \
    --line-bytes=${uncompressed_chunk_size} \
    --filter="$compressor | $writer" \
    - ${filename}.part.