#!/bin/bash

# Basic config
OUTPUT_DIRECTORY=${OUTPUT_DIRECTORY:-'./output'}

# Change this if you want to specify a path to use a specific version of FFMPeg
FFMPEG=${FFMPEG:-'ffmpeg'}

# Number of threads which will be used for transcoding. With newer FFMPEGs and x264
# encoders "0" means "optimal". This is normally the number of CPU cores.
NUMTHREADS=${NUMTHREADS:-"0"}

# Video codec for the output video. Will be used as an value for the -vcodec argument
VIDEO_CODEC=${VIDEO_CODEC:-"libx264"}

# Audio codec for the output video. Will be used as an value for the -acodec argument
AUDIO_CODEC=${AUDIO_CODEC:-"aac"}

# Additional flags for ffmpeg
FFMPEG_FLAGS=${FFMPEG_FLAGS:-""}

# If the input is a live stream (i.e. linear video) this should be 1
LIVE_STREAM=${LIVE_STREAM:-0}

# Video bitrates to use in output (comma seperated list if you want to create an adaptive stream.)
# leave null to use the input bitrate
OP_BITRATES=${OP_BITRATES:-''}

# Determines whether the processing for adaptive streams should run sequentially or not
NO_FORK=${NO_FORK:-0}

# Path Prefix (Not incl. Key)
PATH_PREFIX=${PATH_PREFIX:-''}

# Key Path Prefix
KEY_PREFIX=${KEY_PREFIX:-''}

# Lets put our functions here


## Output the script's CLI Usage
#
#
function print_usage(){

cat << EOM
Usage: c2hls -[lf] [-c segmentcount] -i [inputfile] -s [segmentlength(seconds)] -o [outputdir] -b [bitrates]  [-p filename]
	-i	Input file
	-s	Segment length (seconds)
	-o	Output directory (default: ./output)
	-l	Input is a live stream
	-c	Number of segments to include in playlist (live streams only) - 0 is no limit
	-b	Output video Bitrates (comma seperated list for adaptive streams)
	-f	Foreground encoding only (don't fork the encoding processes into the background - adaptive non-live streams only)
	-p	Playlist filename
	-t	Segment filename prefix
	-S	Segment directory name (default none)
	-e	Encrypt the HLS segments (default none)
	-2	2-pass encoding
	-q	Quality (changes to CRF)
	-C	Constant Bit Rate (CBR as opposed to AVB)
	-u	Path Prefix
	-k	Key Prefix
	-K	Key Name
EOM

exit

}


function createStream(){
# For VoD and single bitrate streams the variables we need will exist in Global scope.
# for live adaptive streams though, that won't be the case, so we need to take them as arguments
# Some are global though, so we'll leave those as is

playlist_name="$1"
output_name="$2"
bitrate="$3"
infile="$4"


# Resolution comes from global $resolution
#
# See HLS-27
#
# I don't like this use of $resolution when it's not local, but there was already a use 
# of non-local $br in here so I'll fix this later


local PASSVAR=
if $TWOPASS; then
	local LOGFILE="$OUTPUT_DIRECTORY/bitrate$br"
	PASSVAR="-passlogfile \"$LOGFILE\" -pass 2"

	$FFMPEG -i "$infile" \
		-pass 1 \
		-passlogfile "$LOGFILE" \
		-an \
		-vcodec libx264 \
		-f mpegts \
		$bitrate \
		$resolution \
		$FFMPEG_ADDITIONAL \
		-loglevel error -y \
		/dev/null
fi

$FFMPEG -i "$infile" \
    $PASSVAR \
    -y \
    -vcodec "$VIDEO_CODEC" \
    -acodec "$AUDIO_CODEC" \
    -threads "$NUMTHREADS" \
    -map 0 \
    -flags \
    -global_header \
    -f segment \
    -segment_list "$playlist_name" \
    -segment_time "$SEGLENGTH" \
    -segment_format mpeg_ts \
    $resolution \
    $bitrate \
    $FFMPEG_ADDITIONAL \
    $FFMPEG_FLAGS \
    "$OUTPUT_DIRECTORY/$output_name"
}


function createVariantPlaylist(){
playlist_name="$1"
echo "#EXTM3U" > "$playlist_name"
}


function appendVariantPlaylistentry(){
playlist_name=$1
playlist_path=$2
bw_statement=$3
m3u8_resolution=''

if [[ "$bw_statement" == *"-"* ]]
then
    m3u8_resolution=",RESOLUTION=$(echo "$bw_statement" | cut -d- -f2)"
    bw_statement=$(echo "$bw_statement" | cut -d- -f1) 
fi

playlist_bw=$(( $bw_statement * 1000 )) # bits not bytes :)

cat << EOM >> "$playlist_name"
#EXT-X-STREAM-INF:BANDWIDTH=$playlist_bw$m3u8_resolution
$PATH_PREFIX$playlist_path
EOM

}


function awaitCompletion(){
# Monitor the encoding pids for their completion status
while [ ${#PIDS[@]} -ne 0 ]; do
    # Calculate the length of the array
    pid_length=$((${#PIDS[@]} - 1))

    # Check each PID in the array
    for i in `seq 0 $pid_length`
    do
	  # Test whether the pid is still active
	  if ! kill -0 ${PIDS[$i]} 2> /dev/null
	  then
		echo "Encoding for bitrate ${BITRATE_PROCESSES[$i]}k completed"

		if [ "$LIVE_STREAM" == "1" ] && [ `grep 'EXT-X-ENDLIST' "$OUTPUT_DIRECTORY/${PLAYLIST_PREFIX}_${BITRATE_PROCESSES[$i]}.m3u8" | wc -l ` == "0" ]
		then
		    # Correctly terminate the manifest. See HLS-15 for info on why
		    echo "#EXT-X-ENDLIST" >> "$OUTPUT_DIRECTORY/${PLAYLIST_PREFIX}_${BITRATE_PROCESSES[$i]}.m3u8"
		fi

		unset BITRATE_PROCESSES[$i]
		unset PIDS[$i]
	  fi
    done
    PIDS=("${PIDS[@]}") # remove any nulls
    BITRATE_PROCESSES=("${BITRATE_PROCESSES[@]}") # remove any nulls
    sleep 1
done
}

function encrypt(){
# Encrypt the generated segments with AES-128 bits


    # Only run the encryption routine if it's been enabled  (and not blocked)
    if [ ! "$ENCRYPT" == "1" ] || [ "$LIVE_STREAM" == "1" ]
    then
        return
    fi

    echo "Generating Encryption Key"
    KEY_FILE="$OUTPUT_DIRECTORY/${KEY_NAME}.key"

    openssl rand 16 > $KEY_FILE
    ENCRYPTION_KEY=$(cat $KEY_FILE | hexdump -e '16/1 "%02x"')

    echo "Encrypting Segments"
    for SEGMENT_FILE in ${OUTPUT_DIRECTORY}/*.ts
    do
        SEG_NO=$( echo "$SEGMENT_FILE" | grep -o -P '_[0-9]+\.ts' | tr -dc '0-9' )
        ENC_FILENAME="$OUTPUT_DIRECTORY/${SEGMENT_PREFIX}_enc_${SEG_NO}".ts

        # Strip leading 0's so printf doesn't think it's octal
        #SEG_NO=${SEG_NO##+(0)} # Doesn't work for some reason - need to check shopt to look further into it
        SEG_NO=$(echo $SEG_NO | sed 's/^0*//' )
        
        # Convert the segment number to an IV. 
	INIT_VECTOR=$(printf '%032x' $SEG_NO)
	openssl aes-128-cbc -e -in $SEGMENT_FILE -out $ENC_FILENAME -nosalt -iv $INIT_VECTOR -K $ENCRYPTION_KEY

        # Move encrypted file to the original filename, so that the m3u8 file does not have to be changed
        mv $ENC_FILENAME $SEGMENT_FILE
        
    done

    echo "Updating Manifests"
    # this isn't technically correct as we needn't write into the master, but should still work
    for manifest in ${OUTPUT_DIRECTORY}/*.m3u8
    do
        # Insert the KEY at the 5'th line in the m3u8 file
        sed -i "5i #EXT-X-KEY:METHOD=AES-128,URI="${KEY_PREFIX}${KEY_NAME}.key "$manifest"
    done
}

# This is used internally, if the user wants to specify their own flags they should be
# setting FFMPEG_FLAGS
FFMPEG_ADDITIONAL=''
LIVE_SEGMENT_COUNT=0
IS_FIFO=0
TMPDIR=${TMPDIR:-"/tmp"}
MYPID=$$
TWOPASS=false
QUALITY=
CONSTANT=false
# Get the input data

# This exists to maintain b/c
LEGACY_ARGS=1

# If even one argument is supplied, switch off legacy argument style
while getopts "i:o:s:c:b:p:t:S:q:u:k:K:Clfe2" flag
do
	LEGACY_ARGS=0
        case "$flag" in
                i) INPUTFILE="$OPTARG";;
                o) OUTPUT_DIRECTORY="$OPTARG";;
                s) SEGLENGTH="$OPTARG";;
		l) LIVE_STREAM=1;;
		c) LIVE_SEGMENT_COUNT="$OPTARG";;
		b) OP_BITRATES="$OPTARG";;
		f) NO_FORK=1;;
		p) PLAYLIST_PREFIX="$OPTARG";;
		t) SEGMENT_PREFIX="$OPTARG";;
		S) SEGMENT_DIRECTORY="$OPTARG";;
		e) ENCRYPT=1;;
		2) TWOPASS=true;;
		q) QUALITY="$OPTARG";;
		C) CONSTANT=true;;
		u) PATH_PREFIX="$OPTARG";;
		k) KEY_PREFIX="$OPTARG";;
		K) KEY_NAME="$OPTARG";;
        esac
done


if [ "$LEGACY_ARGS" == "1" ]
then
  # Old Basic Usage is 
  # cmd.sh inputfile segmentlength 

  INPUTFILE=${INPUTFILE:-$1}
  SEGLENGTH=${SEGLENGTH:-$2}
  if ! [ -z "$3" ]
  then
    OUTPUT_DIRECTORY=$3
  fi
fi


# Check we've got the arguments we need
if [ "$INPUTFILE" == "" ] || [ "$SEGLENGTH" == "" ]
then
  print_usage
fi

# FFMpeg is a pre-requisite, so let check for it
if hash $FFMPEG 2> /dev/null
then
  # FFMpeg exists
  echo "ffmpeg command found.... continuing"
else
  # FFMPeg doesn't exist, uh-oh!
  echo "Error: FFmpeg doesn't appear to exist in your PATH. Please addresss and try again"
  exit 1
fi

# Check whether the input is a named pipe
if [ -p "$INPUTFILE" ]
then
  echo "Warning: Input is FIFO - EXPERIMENTAL"
  IS_FIFO=1
fi

# Make sure that the trailing slashes are added
if [ "$PATH_PREFIX" != "" ] && [ "${PATH_PREFIX:$length:1}" != "/" ]
then
  PATH_PREFIX+="/"
fi

if [ "$KEY_PREFIX" != "" ] && [ "${KEY_PREFIX:$length:1}" != "/" ]
then
  KEY_PREFIX+="/"
fi

# Check output directory exists otherwise create it
if [ ! -w $OUTPUT_DIRECTORY ]
then
  echo "Creating $OUTPUT_DIRECTORY"
  mkdir -p $OUTPUT_DIRECTORY
fi

if [ "$LIVE_STREAM" == "1" ]
then
    FFMPEG_ADDITIONAL+="-segment_list_flags +live"

    if [ "$LIVE_SEGMENT_COUNT" -gt 0 ]
    then
	WRAP_POINT=$(($LIVE_SEGMENT_COUNT * 2)) # Wrap the segment numbering after 2 manifest lengths - prevents disks from filling
	FFMPEG_ADDITIONAL+=" -segment_list_size $LIVE_SEGMENT_COUNT -segment_wrap $WRAP_POINT"
    fi
fi


# Pulls file name from INPUTFILE which may be an absolute or relative path.
INPUTFILENAME=${INPUTFILE##*/}

# If a prefix hasn't been specified, use the input filename
PLAYLIST_PREFIX=${PLAYLIST_PREFIX:-$INPUTFILENAME}
SEGMENT_PREFIX=${SEGMENT_PREFIX:-$PLAYLIST_PREFIX}
KEY_NAME=${KEY_NAME:-$PLAYLIST_PREFIX}

# The 'S' option allows segments and bitrate specific manifests to be placed in a subdir
SEGMENT_DIRECTORY=${SEGMENT_DIRECTORY:-''}

if [ ! "$SEGMENT_DIRECTORY" == "" ]
then

	if [ ! -d "${OUTPUT_DIRECTORY}/${SEGMENT_DIRECTORY}" ]
	then
		mkdir "${OUTPUT_DIRECTORY}/${SEGMENT_DIRECTORY}"
	fi

	SEGMENT_DIRECTORY+="/"
	OUTPUT_DIRECTORY+="/"
fi

# Set the bitrate
if [ ! "$OP_BITRATES" == "" ]
then
      # Make the bitrate list easier to parse
      OP_BITRATES=${OP_BITRATES//,/$'\n'}

      # Create an array to house the pids for backgrounded tasks
      declare -a PIDS
      declare -a BITRATE_PROCESSES

      # Get the variant playlist created
      createVariantPlaylist "$OUTPUT_DIRECTORY/${PLAYLIST_PREFIX}_master.m3u8"
      for br in $OP_BITRATES
      do
            bw=$br
            if [[ "$br" == *"-"* ]]
            then
                bw=$(echo "$br" | cut -d- -f1) 
            fi
      
	    appendVariantPlaylistentry "$OUTPUT_DIRECTORY/${PLAYLIST_PREFIX}_master.m3u8" "${SEGMENT_DIRECTORY}${PLAYLIST_PREFIX}_${bw}.m3u8" "$br"
      done

      OUTPUT_DIRECTORY+=$SEGMENT_DIRECTORY

      # Now for the longer running bit, transcode the video
      for br in $OP_BITRATES
      do
      
              # Check whether there's a resolution included in the bitrate string
              #
              # See HLS-27
              if [[ "$br" == *"-"* ]]
              then
                resolution="-vf scale=$(echo "$br" | cut -d- -f2 | sed 's/x/:/')"
                br=$(echo "$br" | cut -d- -f1) 
              fi
      
              if [ -z $QUALITY ]; then
		if $CONSTANT; then
	          BITRATE="-b:v ${br}k -bufsize ${br}k -minrate ${br}k -maxrate ${br}k"
		else
	          BITRATE="-b:v ${br}k"
		fi
	      else
	        BITRATE="-crf $QUALITY -maxrate ${br}k -bufsize ${br}k"
                if [ $VIDEO_CODEC = "libx265" ]; then
                  BITRATE="$BITRATE -x265-params --vbv-maxrate ${br}k --vbv-bufsize ${br}k"
                fi
	      fi
	      echo "Bitrate options: $BITRATE"
	      # Finally, lets build the output filename format
	      OUT_NAME="${SEGMENT_PREFIX}_${br}_%05d.ts"
	      PLAYLIST_NAME="$OUTPUT_DIRECTORY/${PLAYLIST_PREFIX}_${br}.m3u8"
	      SOURCE_FILE="$INPUTFILE"
	      echo "Generating HLS segments for bitrate ${br}k - this may take some time"

	      if [ "$NO_FORK" == "0" ] || [ "$LIVE_STREAM" == "1" ]
	      then
		      # Processing Starts
		      if [ "$IS_FIFO" == "1" ]
		      then
			    # Create a FIFO specially for this bitrate
			    SOURCE_FILE="$TMPDIR/hlsc.encode.$MYPID.$br"
			    mknod "$SOURCE_FILE" p
		      fi

		      # Schedule the encode
		      createStream "$PLAYLIST_NAME" "$OUT_NAME" "$BITRATE" "$SOURCE_FILE" &
		      PID=$!
		      PIDS=(${PIDS[@]} $PID)
		      BITRATE_PROCESSES=(${BITRATE_PROCESSES[@]} $br)
	      else
		      createStream "$PLAYLIST_NAME" "$OUT_NAME" "$BITRATE" "$SOURCE_FILE"
	      fi

      done

      if [ "$IS_FIFO" == "1" ]
      then
	      # If the input was a FIFO we need to read from it and push into the new FIFOs
	      cat "$INPUTFILE" | tee $(for br in $OP_BITRATES; do echo "$TMPDIR/hlsc.encode.$MYPID.$br"; done) > /dev/null &
	      TEE_PID=$!
      fi

      if [ "$NO_FORK" == "0" ] || [ "$LIVE_STREAM" == "1" ]
      then
	    # Monitor the background tasks for completion
	    echo "All transcoding processes started, awaiting completion"
	    awaitCompletion
	    
	    # As of HLS-20 encrypt will only run if the relevant vars are set
	    encrypt
      fi

      if [ "$IS_FIFO" == "1" ]
      then
	    for br in $OP_BITRATES
	    do 
		rm -f "$TMPDIR/hlsc.encode.$MYPID.$br"; 
	    done
	    # If we were interrupted, tee may still be running
	    kill $TEE_PID 2> /dev/null 
      fi

else

  OUTPUT_DIRECTORY+=$SEGMENT_DIRECTORY
  # No bitrate specified

  # Finally, lets build the output filename format
  OUT_NAME="${SEGMENT_PREFIX}_%05d.ts"
  PLAYLIST_NAME="$OUTPUT_DIRECTORY/${PLAYLIST_PREFIX}.m3u8"

  echo "Generating HLS segments - this may take some time"

  # Processing Starts

  createStream "$PLAYLIST_NAME" "$OUT_NAME" "$BITRATE" "$INPUTFILE"

  # As of HLS-20 encrypt will only run if the relevant vars are set
  encrypt
fi
