#!/bin/bash
#
# https://ocv.me/ffmpeg/webm.sh
# ed <irc.rizon.net>, MIT-licensed, 2018-03-12 03:46
#
# takes a media file and spits out webm or mp3 files
# subtitles are burned into the video
# volume is normalized to -12dB RMS
# ffmpeg colorspace bug is compensated for
# supports OSX and Linux (glibc and musl)
# 
set -e


###
### OPTIONS
###
	# vp8 is close to 7x as fast as vp9
	# so we'll force vp8 until -row-mt is added
	force_vp8='yes'
	force_vp9=''

	# if no encoder is forced above, default to vp8
	# but change to vp9 if there is more than 4 cpu cores
	vp9_threshold=4
###
### END OF OPTIONS
###


# setting color_range without these makes it go N/A
# and setting color_range BEFORE these /also/ breaks
hey_another_bug="-color_primaries bt709 -color_trc bt709 -colorspace bt709"

# make ffmpeg as silent as sensible
silent="-hide_banner -v error"
chatty="-hide_banner"


function need()
{
	for p in "$@"
	do
		which $p |
		grep -E ... &&
			return ||
			true
	done
	[[ "x$2" == "x" ]] &&
	{
		printf '\033[1;31m\n   need this program:  %s\033[0m\n\n' "$1" >&2
		exit 1
	}
	printf '\033[1;31m\n   need either of these programs:  %s\033[0m\n\n' "$*" >&2
	exit 1
}

need awk > /dev/null
need ffmpeg > /dev/null
need ffprobe > /dev/null
player="$(need ffplay mpv)"

player_args=''
printf '%s\n' "$player" |
grep -qE 'ffplay[^/]*$' &&
	player_args='-autoexit -probesize 32'

# osx support
echo a | sed -r 's/a/b/' 2>/dev/null | grep -qE '^b$' ||
{
	function sed() { gsed "$@"; }
	function find() { gfind "$@"; }

	ffmpeg -codecs 2>&1 | grep -qE 'encoders:.* libvorbis' &&
	ffmpeg -codecs 2>&1 | grep -qE 'encoders:.* libvpx[^-]' ||
	{
		printf '\033[1;31m\n\n  your ffmpeg is useless, run this:\n\033[0;33m    '
		printf 'brew reinstall --force ffmpeg --with-tools --with-sdl2 --with-freetype --with-libass --with-libvorbis --with-libvpx --with-fdk-aac --with-opus\n\n\033[0m'
		printf "  (fdk-aac is not really necessary but you'll thank me later)\n\n"
		exit 1
	}
}

function wt() {
	printf '\033]0;%s\007' "$*"
}

function cfg2human()
{
	expect="$1"
	cfg="$2"
	ffmpeg $silent \
		$hey_another_bug $cfg \
		-f lavfi -i smptebars -frames:v 1 \
		-vcodec libx264 -preset ultrafast \
		$hey_another_bug $cfg \
		-f nut - </dev/null |
	
	ffprobe $silent -show_streams - |
	grep -E "^$expect="
}

grep -qE '^EOF$' webm.cache 2>/dev/null ||
{
	tmp="$(mktemp)"
	printf '\n# \033[1;33m%-32s %s\033[0m\n' "this is what ffprobe says" "when you make a video with" > webm.cache
	printf '\ngenerating profile cache...\n\n' >&2
	ffmpeg $silent -h full </dev/null | awk '
		
		/^  -/ {
			k = ""
		}
		
		k != "" {
			v1 = $0
			v2 = $0
			sub(/^ */, "", v1)
			sub(/ .*/, "", v1)
			printf("%s %s\n", k, v1)
		}
		
		/^  -colorspace .*default/ {
			k = "1 color_space -colorspace"
		}

		/^  -color_trc .*default/ {
			k = "2 color_transfer -color_trc"
		}
		
		/^  -color_primaries .*default/ {
			k = "3 color_primaries -color_primaries"
		}
		
		/^  -color_range .*default/ {
			k = "4 color_range -color_range"
		}

	' |
	grep -vE ' unspecified' > "$tmp"

	tasks=$(cat "$tmp" | wc -l)
	
	while IFS=' ' read -r order expect cfg
	do
		printf "\033[A%3s " $tasks >&2
		printf "%-${tasks}s\n" "" |
		sed -r 's/ /#/g;s/$/./' >&2
		tasks=$((tasks-1))
		
		result="$(cfg2human "$expect" "$cfg")" &&
		printf '%s\n' "$result" | grep -qE '=' &&
		printf '%s %-32s %s\n' "$order" "$result" "$cfg"
	done < <(cat "$tmp") |
	grep -vE '=unknown' |
	sort >> webm.cache
	
	echo EOF >> webm.cache
	printf '\033[Acache ok\033[J\n\n' >&2
	rm "$tmp"
}

function gencsp()
{
	[[ "x$1" == "x" ]] &&
	{
		printf '\n\033[1;31m   need argument 1:  video file\033[0m\n\n'
		exit 1
	}

	ffprobe $silent -show_streams "$1" </dev/null |
	grep -E '^color_' |
	head -n 4 |
	while IFS= read -r ln
	do
		grep -E "^. ${ln} " webm.cache | head -n 1 || true
	done |
	sort |
	while read junk junk v
	do
		printf '%s ' "$v"
	done
	echo
}

function h2sec()
{
	rv=0
	iv="$1"
	
	printf '%s\n' "$iv" | grep -qE ':.*:' &&
	{
		seg="$(printf '%s\n' "$iv" | sed -r 's/:.*//')"
		iv="$(printf '%s\n' "$iv" | sed -r 's/[^:]*://;s/^0*$//;s/^$/0/')"
		rv=$((rv+seg))
	}
	rv=$((rv*60))
	
	printf '%s\n' "$iv" | grep -qE ':' &&
	{
		seg="$(printf '%s\n' "$iv" | sed -r 's/:.*//')"
		iv="$(printf '%s\n' "$iv" | sed -r 's/[^:]*://;s/^0*$//;s/^$/0/')"
		rv=$((rv+seg))
	}
	rv=$((rv*60))
	
	iv="$(
		printf '%s\n' "$iv" |
		sed -r '
			s/\./,/g;
			s/^([^,]*)$/\1,/;
			s/$/000/;
			s/(.*,...).*/\1/
	')"
	int="$(printf '%s\n' "$iv" | sed -r 's/,.*//')"
	frac="$(printf '%s\n' "$iv" | sed -r 's/.*,//')"
	rv=$((rv+int))
	printf '%s.%s\n' "$rv" "$frac"
}

function sub()
{
	i1=$(printf '%s\n' "$1" | sed -r 's/\..*//;s/^0*//;s/^$/0/')
	i2=$(printf '%s\n' "$2" | sed -r 's/\..*//;s/^0*//;s/^$/0/')
	f1=$(printf '%s\n' "$1" | sed -r 's/^[^\.]*$//;s/.*\.//;s/$/000/;s/(...).*/\1/;s/^0*//;s/^$/0/')
	f2=$(printf '%s\n' "$2" | sed -r 's/^[^\.]*$//;s/.*\.//;s/$/000/;s/(...).*/\1/;s/^0*//;s/^$/0/')
	#printf '[%s] [%s] [%s] [%s] [%s]\n' $i1 $i2 $f1 $f2 x >&2
	fv=0
	iret=$((i1-i2))
	[[ $f2 -gt $f1 ]] &&
	{
		iret=$((iret-1))
		f1=$((f1+1000))
	}
	fv=$((f1-f2))
	printf '%s.000%s\n' "$iret" "$fv" |
		sed -r 's/\.0*(...)$/.\1/'
}

#sub $1 $2
#exit 0

# ./webm.sh subfull.mkv 1 1 640 6.9 96.8 1

infile="$1"
atrack="$2"
strack="$3"
xres="$4"
yres="$5"
t1=$(h2sec "$6")
t2=$(h2sec "$7")
td=$(sub $t2 $t1)
outfile="$8"

mp3=''
webm=''
mp3="$( printf '%s\n' "$outfile" | grep -qE '\.mp3$'  && echo 1 || true)"
webm="$(printf '%s\n' "$outfile" | grep -qE '\.webm$' && echo 1 || true)"
atrack=$((atrack-1))
strack=$((strack-1))

printf '\033[36minput file:      \033[0m%s\n' "$infile"
printf '\033[36maudio track:     \033[0m%s\n' "$((atrack+1))"
printf '\033[36msubtitle track:  \033[0m%s\n' "$((strack+1))"
printf '\033[36mhorizontal res:  \033[0m%s\n' "$xres"
printf '\033[36mvertical res:    \033[0m%s\n' "$yres"
printf '\033[36mlower timestamp: \033[0m%s\n' "$t1"
printf '\033[36mupper timestamp: \033[0m%s\n' "$t2"
printf '\033[36moutput duration: \033[0m%s\n' "$td"
printf '\033[36moutput to:       \033[0m%s\n' "$outfile"
printf '\n'

[[ "x$7" == "x" ]] &&
{
	printf '\n'
	printf '   need argument 1:  \033[36m[filepath]\033[33m input file\n\033[0m'
	printf '   need argument 2:  \033[36m[a number]\033[33m audio track\n\033[0m'
	printf '   need argument 3:  \033[36m[a number]\033[33m subtitle track\n\033[0m'
	printf '   need argument 4:  \033[36m[a number]\033[33m max horizontal resolution\n\033[0m'
	printf '   need argument 5:  \033[36m[a number]\033[33m max vertical resolution\n\033[0m'
	printf '   need argument 6:  \033[36m[HH:MM:SS]\033[33m start position\n\033[0m'
	printf '   need argument 7:  \033[36m[HH:MM:SS]\033[33m end position\n\033[0m'
	printf '   need argument 8:  \033[33mtestmode or output file\n\033[0m'
	printf '\033[36m\n'
	printf '   1 is the first audio track, 2 is the second, etc.\n'
	printf '   1 is the first subtitle track, 2 is the second, etc.\n'
	printf '   disable audio or subtitles with 0\n'
	printf '\033[0m\n'
	printf '   video will never be wider than max horizontal res,\n'
	printf '                   or taller than max vertical res\n'
	printf '\033[0m\n'
	printf '   if arg.8 is the letter s:   \033[33mshow start of segment\n\033[0m'
	printf '   if arg.8 is the letter e:   \033[33mshow end of segment\n\033[0m'
	printf '   if arg.8 is something.mp3:  \033[33mencode just audio to this file\n\033[0m'
	printf '   if arg.8 is something.webm: \033[33mencode video+audio to this file\n\033[0m'
	printf '\n'
	printf '\n'
	printf 'webm example:\n'
	printf './webm.sh \033[1;31mrailgun.mkv \033[1;33m1   \033[1;32m1  \033[1;36m1280  \033[1;34m720 \033[1;35m02:24.13 \033[1;31m03:53.72 \033[1;33mrailgun-op.webm\n'
	printf '          \033[1;31m\\inputfile/ \033[1;33mau \033[1;32msub \033[1;36mxres \033[1;34myres \033[1;35m\\start_/ \033[1;31m\\_end__/ \033[1;33moutput filename\033[0m\n'
	printf '\n'
	printf 'mp3 example:\n'
	printf './webm.sh \033[1;31mrailgun.mkv \033[1;33m1   \033[1;32m0  \033[1;34mx x \033[1;35m02:24.13 \033[1;31m03:53.72 \033[1;33mrailgun-op.mp3\n'
	printf '          \033[1;31m\\inputfile/ \033[1;33mau \033[1;32msub \033[1;34mres \033[1;35m\\start_/ \033[1;31m\\_end__/ \033[1;33moutput filename\n'
	printf '\033[0m\n'
	exit 1
}

#[[ $mp3 ]] && echo mp3 || echo webm
#exit 0

# see if ffmpeg is new enough to increase missing b-frame tolerance
skip_orphaned_bframes='-max_muxing_queue_size 99999'
ffmpeg -h full | grep -q max_muxing_queue_size ||
	skip_orphaned_bframes=''

naudio=$(
	ffprobe $silent -show_streams "$infile" |
	grep -E '^codec_type=audio$' | wc -l
)
[[ $naudio -eq 0 ]] &&
{
	printf '\033[1;33mthere is no audio in your file; disabling audio\n\033[0m'
	atrack=-1
}
[[ $naudio -le $atrack ]] &&
{
	printf '\n  \033[1;31myou asked for audio %s but there are only %s audio tracks\n\n\033[0m' $((atrack+1)) $naudio
	exit 1
}

nsubs=$(
	ffprobe $silent -show_streams "$infile" |
	grep -E '^codec_type=subtitle$' | wc -l
)
[[ $nsubs -eq 0 ]] &&
{
	printf '\033[1;33mthere are no subtitles in your file; disabling subtitles\n\033[0m'
	strack=-1
}
[[ $nsubs -le $strack ]] &&
{
	printf '\n  \033[1;31myou asked for subtitle %s but there are only %s subtitles\n\n\033[0m' $((strack+1)) $nsubs
	exit 1
}

printf '\033[36m * seeking to %s, reading until %s, duration %s\033[0m\n' $t1 $t2 $td

which lscpu >/dev/null &&
	ncores=$(lscpu -p | grep -E '^[0-9]' | wc -l) ||
	ncores=$(sysctl -n hw.ncpu)

[[ $ncores -gt $vp9_threshold ]] &&
	vcodec=vp9 ||
	vcodec=vp8

[[ $force_vp8 ]] && vcodec=vp8
[[ $force_vp9 ]] && vcodec=vp9

printf '\033[36m * using %s cores to encode, will create %s media\033[0m\n' $ncores $vcodec
printf '\n'

live=''
[[ "x$outfile" == "xs" ]] && live=1
[[ "x$outfile" == "xe" ]] && live=1 && t1=$(sub $t2 1)

subfilter="setpts=PTS+${t1}/TB,subtitles=${infile}:si=${strack},setpts=PTS-STARTPTS,"
[[ $strack -lt 0 ]] &&
	subfilter=''

# compare against 16/10 since ntsc is always shorter than pal
scalefilter="scale='if(gt(a*sar,16/10),${xres},-4)':'if(gt(a*sar,16/10),-8,${yres})'"

csp="$(gencsp "$infile")"

[[ $live ]] &&
{
	ffmpeg $silent -y $csp -vsync vfr -ss $t1 -t 1 -i "$infile" \
		-map 0:v:0 -map 0:a:${atrack}? -sn \
		-map_metadata -1 -map_chapters -1 \
		-vf "${subfilter}${scalefilter}" $skip_orphaned_bframes \
		-vcodec rawvideo -acodec pcm_s16le -f nut $csp - < /dev/null |
	"$player" $player_args -
	exit $?
}

[[ $mp3 ]] || [[ $webm ]] ||
{
	printf '\033[1;31m\n  filename must be either "s", "e", some.mp3 or some.webm\033[0m\n\n'
	exit 1
}

function volseg_v1()
{
	cat "$1" |
	grep -E "^.\"${2}\" : \"" |
	sed -r 's/.*"([^"]*)",?$/\1/'
}

function genvolfilter_v1()
{
	tmp=$(mktemp)
	
	# ffmpeg -i in.wav -af loudnorm=I=-16:TP=-1.5:LRA=11:print_format=json -f null - 
	# ffmpeg -i in.wav -af loudnorm=I=-16:TP=-1.5:LRA=11:measured_I=-27.61:measured_LRA=18.06:measured_TP=-4.47:measured_thresh=-39.20:offset=0.58:linear=true:print_format=summary -ar 48k out.wav
	# "input_i" : "-27.61",
	# "input_lra" : "18.06",
	# "input_tp" : "-4.47",
	# "input_thresh" : "-39.20",
	# "target_offset" : "0.58"
	
	af_base="loudnorm=I=-12:TP=-1.5:LRA=11"
	
	ffmpeg $chatty -ss $t1 -i "$infile" -t $td -map 0:a:${atrack} \
		-af "$af_base":print_format=json \
		-f null - < /dev/null > "$tmp" 2>&1 ||
	{
		printf '\n\033[1;31merror getting volume info:\n' >&2
		cat "$tmp" >&2
		printf '\033[0m\n' >&2
		rm "$tmp"
		exit 1
	};
	
	#printf '%s\n' "$tmp"
	#exit 0
	
	printf '%s:measured_I=%s:measured_LRA=%s:measured_TP=%s:measured_thresh=%s:offset=%s:linear=true:print_format=summary\n' \
		"$af_base" \
		"$(volseg_v1 "$tmp" input_i)" \
		"$(volseg_v1 "$tmp" input_lra)" \
		"$(volseg_v1 "$tmp" input_tp)" \
		"$(volseg_v1 "$tmp" input_thresh)" \
		"$(volseg_v1 "$tmp" target_offset)"
	
	rm "$tmp"
}

function volseg()
{
	cat "$1" |
	grep -E '\[Parsed_volumedetect_0 @ 0x' |
	sed -r 's/.*\] //' |
	grep -E "^$2:" |
	awk '{print $2}'
}

function intify()
{
	printf '%s\n' "$1" |
	sed -r '
		s/^([^\.]*)$/\1./;
		s/$/000/;
		s/\.(...).*/\1/;
		s/\.//'
}

function floatify()
{
	printf '%s\n' "$1" |
	sed -r 's/(.{3})$/.\1/'
}

function genvolfilter()
{
	tmp=$(mktemp)
	
	af_base="loudnorm=I=-${fixme_vol}:TP=-1.5:LRA=11"
	
	ffmpeg $chatty -ss $t1 -i "$infile" -t $td -map 0:a:${atrack} \
		-af "volumedetect" -f null - < /dev/null > "$tmp" 2>&1 ||
	{
		printf '\n\033[1;31merror getting volume info:\n' >&2
		cat "$tmp" >&2
		printf '\033[0m\n' >&2
		rm "$tmp"
		exit 1
	}
	
	mean=$(intify $(volseg "$tmp" mean_volume));
	max=$( intify $(volseg "$tmp" max_volume));
	
	# we'll aim for -12dB, so -12000
	[[ $mean -gt -12000 ]] &&
	{
		# need to provide /some/ filter so let's do this
		printf 'volumedetect'
		rm "$tmp"
		return
	}
	
	gain=$(( (-12000) - (mean) ))
	new_max=$(( max + gain ))
	
	# are we clipping yet?
	[[ $new_max -gt 0 ]] &&
	{
		gain=$(( max * -1 ))
	}
	
	# dump stats
	{
		grep '_volume: ' "$tmp"
		printf 'mean = %s\nmax = %s\ngain = %s\n' $mean $max $(floatify $gain)
		printf '\033[0m'
	} |
	sed -r 's/(.* = )/    \1/' >&2
	
	printf '%s\n' "volume=$(floatify $gain)dB"
	
	rm "$tmp"
}

printf '\033[1;32manalysing volume\033[0m\n' $1
wt "webm-vol-analysis"
[[ $atrack -ge 0 ]] &&
	volfilter="$(genvolfilter)" ||
	volfilter='volumedetect'

passdir="$(mktemp -d)"

function runpass()
{
	printf '\033[1;32mdoing pass %s encoding\033[0m\n' $1
	wt "webm-pass-$1"
	
	[[ "x$1" == "x1" ]] &&
	{
		spd=4
		ftype=null
		ofname=/dev/null
	} ||
	{
		spd=1
		ftype=webm
		ofname="$outfile"
	}
	
	bv=2500k
	mr=4000k
	[[ $xres -lt 1500 ]] &&
	{
		bv=1500k
		mr=2500k
	}

	[[ "x$vcodec" == "xvp8" ]] &&
		vcfg="-vcodec libvpx" ||
		vcfg="-vcodec libvpx-vp9 -tile-columns 6 -frame-parallel 1 -threads $ncores"
	
	ffmpeg \
		$chatty -y $csp -vsync vfr \
		-ss $t1 -t $td -i "$infile" \
		-map 0:v:0 -map 0:a:${atrack}? -sn -map_metadata -1 -map_chapters -1 \
		-vf "${subfilter}${scalefilter}" \
		-af "${volfilter}" -ar 48000 \
		-passlogfile "$passdir/x" \
		-pix_fmt yuv420p $skip_orphaned_bframes \
		-acodec libvorbis -ac 2 -q:a 5 \
		$vcfg -pass $1 -deadline best -speed $spd \
		-minrate 500k -b:v $bv -maxrate $mr -g 200 \
		-lag-in-frames 25 -auto-alt-ref 1 \
		-f $ftype $ofname < /dev/null
}

function makemp3()
{
	printf '\033[1;32mcreating mp3\033[0m\n'
	wt "webm-mp3"
	
	ffmpeg \
		$chatty -y -ss $t1 -t $td -i "$infile" \
		-map 0:a:${atrack} -sn -map_metadata -1 -map_chapters -1 \
		-af "${volfilter}" -ar 48000 \
		-acodec libmp3lame -b:a 192k \
		-f mp3 $outfile < /dev/null
}

[[ $mp3 ]] &&
{
	makemp3
}

[[ $webm ]] &&
{
	runpass 1
	runpass 2
}

rm -rf "$passdir"
printf '\033[1;32m\n  done!\033[0m\n\n'
wt webm-done
