#!/bin/bash

SEG_DURATION_OUTPUT_FILE=seg_duration_output.csv
TRANS_FRAME_OUTPUT_FILE=trans_frame_output.csv
SEG_FRAME_OUTPUT_FILE=seg_frame_output.csv
TRANS_GOP_OUTPUT_FILE=trans_gop_output.csv
SEG_GOP_OUTPUT_FILE=seg_gop_output.csv

###################################

function printUsage() {
  echo
  echo "Usage:"
  echo "      analyze [path/to/transcoded/file]"
  echo 
  echo "This tool analyzes all *.ts_* HLS segments in the local directory for "
  echo "issues. If you pass the path to the transcoded original file as a "
  echo "parameter it will analyze that as well for comparison. Issues that "
  echo "appear in just the HLS analysis point to a miniclient problem. Issues "
  echo "that appear in the latter analysis (or both analyses) point to a "
  echo "BroadCom issue."
  echo
  echo "This tool requires dvbsnoop and ffmpeg!  It also needs gnuplot to "
  echo "print pretty graphs, but it's optional."
  echo
  echo "    'apt-get install dvbsnoop ffmpeg gnuplot'"
  echo
}


## Parse args and check dependencies
FILE=""
if [ $# -eq 1 ]
then
  FILE=$1
fi

RAN_A_TEST=0

## Return true if depedant app (1) exists
dependency_exists() {
    command -v $1 >/dev/null 2>&1
}

## Check for deps without killing terminal (as exit does, hence the goto)
dependency_exists ffmpeg || { echo; echo >&2 "This script requires ffmpeg but it's not installed!"; printUsage; exit; }
dependency_exists dvbsnoop || { echo; echo >&2 "This script requires dvbsnoop but it's not installed!"; printUsage; exit; }


## Gnuplot the results given by the parameters
function plotResult() {
  if dependency_exists gnuplot
  then
    PLOTFILE=$1
    TITLE=$2
    XLABEL=$3
    XRANGE=$4
    YLABEL=$5
    YRANGE=$6

    gnuplot -persist <<HERE
    set title "${TITLE}"
    set xlabel '${XLABEL}'
    set xrange[${XRANGE}]
    set ylabel '${YLABEL}'
    set yrange[${YRANGE}]
    unset key
    plot "${PLOTFILE}" using 1:3 with lines
HERE
  else
    echo "   'gnuplot' is not installed, so no pretty graphs for you today :("
  fi
}


## Calculate the GOP between each keyframe given an input file (1) to the given output file (2)
function calulateGOP() {
  OUTPUT_FILE=$2

  ## GOP - Algorithm stolen Shamelessly From Zeev Leiber
  PID=$(ffprobe -i $1 2>&1 | grep "Video:" | sed 's/[][]/ /g' | awk '{print "echo $((" $3 "))"}' | bash)
  echo "   Processing $1 (Video PID $PID)..." 
  
  GOPs=$(dvbsnoop -s ts -if $1 $PID | grep random_access | awk '{print $2}')
  FRAME_COUNT=0
  for g in $GOPs
  do
    if [[ $g -eq 1 ]]
    then
      if [[ $FRAME_COUNT -gt 0 ]]
      then
        echo $NUM_GOP	$1	$FRAME_COUNT >> "$OUTPUT_FILE"
        FRAME_COUNT=1
        NUM_GOP=`expr $NUM_GOP + 1`
      fi
    else 
      FRAME_COUNT=`expr $FRAME_COUNT + 1`
    fi
  done
}


## Calculate the size of each frame given an input file (1) adding results into a single file (2)
function calculateFrameSize() {
  OUTPUT_FILE=$2
  echo "   Processing $1..." 

  FRAMEs=$(ffprobe -v quiet -show_frames -select_streams v -i $1 | awk '/pkt_size=/ {sub(/pkt_size=/, "", $1); print $1;}')
  for g in $FRAMEs
  do
    echo $NUM_FRAMES	$1	$g >> $OUTPUT_FILE
    NUM_FRAMES=`expr $NUM_FRAMES + 1`
  done
}


#####
## Work around fact that you can't pass multi-line input to a function...

## Calculate GOP for each segment adding results into a single file (1)
function calculateSegmentGOPs() {
  FILES=$(ls -v *.ts_*)

  echo "# NUM	File	GOP" > $1
  NUM_GOP=1
  for f in $FILES
  do
  	calulateGOP $f $1
  done

  echo
  echo "Result written to $1"

  plotResult $1 "HLS GOP Analysis" "Group" "1:$NUM_GOP-1" "GOP" "0:35"
}

## Calculate GOP for transcoded file (1) adding results into a single file (2)
function calculateTranscodedGOP() {
  echo "# NUM	File	GOP" > $2

  NUM_GOP=1
  calulateGOP $1 $2

  echo
  echo "Result written to $2"

  plotResult $2 "Transcoded GOP Analysis" "Group" "1:$NUM_GOP-1" "GOP" "0:35"
}

## Calculate the size of each frame in each segment adding results into a single file (1)
function calculateSegmentFS() {
  FILES=$(ls -v *.ts_*)

  echo "# NUM	File	Frame_Size" > $1
  NUM_FRAMES=1
  for f in $FILES
  do
  	calculateFrameSize $f $1
  done

  echo
  echo "Result written to $1"

  plotResult $1 "HLS Segment Frame Size Analysis" "Frame" "1:${NUM_FRAMES}-2" "Size (b)" "0:*"
}

## Calculate the size of each frame given an input file (1) adding results into a single file (1)
function calculateTranscodedFS() {
  echo "# NUM	File	Frame_Size" > $2

  NUM_FRAMES=1
  calculateFrameSize $1 $2

  echo
  echo "Result written to $2"

  plotResult $2 "Transcoded Frame Size Analysis" "Frame" "1:${NUM_FRAMES}-2" "Size (b)" "0:*"
}
#####


## Calculate duration of each HLS segment in the current folder
function calculateDurations() {
  echo "# NUM	File	Duration" > $SEG_DURATION_OUTPUT_FILE
  FILES=$(ls -v *.ts_*)
  NUM=1
  for f in $FILES
  do
    echo "   Processing $f..."
    ffprobe -v quiet -show_packets -select_streams v -i $f > ./temp

    duration=$(cat ./temp | awk '/duration_time=0/ {sub(/duration_time=/, "", $1); total = total + $1;}END {print total;}')
    echo $NUM	$f	$duration >> $SEG_DURATION_OUTPUT_FILE
    NUM=`expr $NUM + 1`
  done

  echo
  echo "Result written to $SEG_DURATION_OUTPUT_FILE"

  plotResult $SEG_DURATION_OUTPUT_FILE "HLS Segment Duration" "Segment" "1:${NUM}-1" "Duration (s)" "0:*"
}


## Script flow ##


## Segments
echo
echo "Starting segment analysis (.ts_* files)"
ls -v *.ts_* 2>./temp 1>/dev/null
if [[ $? == 0 ]]
then
  echo
  echo "Analyzing Duration"
  calculateDurations $FILES
  echo
  echo "Analyszing GOP"
  calculateSegmentGOPs $SEG_GOP_OUTPUT_FILE
  echo
  echo "Analyzing Frame Size"
  calculateSegmentFS $SEG_FRAME_OUTPUT_FILE

  RAN_A_TEST=1
else 
  echo "   No segments found, skipping these tests."
fi
echo

## Transcoded
echo "Starting transcoded file analysis"
if [[ "$FILE" != "" ]]
then
  echo
  echo "Analyzing GOP"
  calculateTranscodedGOP $FILE $TRANS_GOP_OUTPUT_FILE
  echo
  echo "Analyzing Frame Size"
  calculateTranscodedFS $FILE $TRANS_FRAME_OUTPUT_FILE

  RAN_A_TEST=1
else
  echo "   No File specified as a parameter, skipping these tests."
fi

## If we didn't test anything, print usage in case the user is new...
if [ $RAN_A_TEST == 0 ]
then
  printUsage
fi

if [ -e "./temp" ]
then
    rm ./temp
fi

echo