#!/bin/bash

set -euo pipefail

VIDEO_CODECS=${VIDEO_CODECS:-"libx264 mpeg4 libvpx-vp9 libvpx"} #ordered by preference

model="realesr-animevideov3"
suffix="anime"
geom="1920x1080"
while (( $# != 0 )) && ( [ "$1" == "-gen" ] || echo "$1" | grep -q "^[0-9]\+x[0-9]\+$" ); do
	if [ "$1" == "-gen" ]; then 
		model="realesrgan-x4plus"
		suffix="general"
	else
		geom="$1"
	fi
	shift
done

if ! which ffprobe >/dev/null || ! which ffmpeg >/dev/null; then
	echo "error: ffprobe / ffmpeg not found"
	exit 1
fi
if ! which mogrify >/dev/null; then
	echo "error: mogrify (ImageMagick) not found"
	exit 1
fi
if ! which realesrgan-ncnn-vulkan >/dev/null; then
	echo "error: realesrgan-ncnn-vulkan not found"
	exit 1
fi

if (( $# == 0 )) || [ "$1" == "--help" ]; then 
	echo "usage: $0 [-gen] [3840×2160] input_video [output_video]"
	exit 2
fi
if [ ! -f "$1" ]; then 
	echo "error: file '$1' not found"
	exit 3
fi
if ! ffprobe "$1" 2>/dev/null; then
	echo "error: '$1' is not a video"
	exit 4
fi

for vcodec in $VIDEO_CODECS; do
	if (( $(ffmpeg -codecs 2>/dev/null | grep " $vcodec " | wc -l) != 0 )); then break; else vcodec=""; fi
done
if [ -z "$vcodec" ]; then
	echo "error: no available video codec (tried: $VIDEO_CODECS)"
	exit 5
fi


inputVideo="$1"
inputFrames=$(echo $(dirname "$inputVideo")/.inputFrames-$$ | sed 's@^\./@@')
outputFrames=$(echo $(dirname "$inputVideo")/.outputFrames-$$ | sed 's@^\./@@')

inputAudioCodec=$(ffprobe -select_streams a:0 -show_entries stream=codec_name "$inputVideo" 2>/dev/null | grep codec_name= | cut -f2 -d=)
if (( $# == 1 )); then
	if ([ "$inputAudioCodec" != "" ] || [ "$inputAudioCodec" != "vorbis" ] || [ "$inputAudioCodec" != "opus" ]) && ([ "$vcodec" == "libvpx-vp9" ] || [ "$vcodec" == "libvpx" ]); then
		extension="webm"
	else #mpeg4
		extension="mp4"
	fi
	outputVideo=$(echo $(dirname "$inputVideo")/$(basename "$inputVideo" | rev | cut -f2- -d. | rev)-$$.$extension | sed 's@^\./@@')
else
	outputVideo="$2"
	extension=$(echo "$outputVideo" | rev | cut -f1 -d. | rev)
	if [ "$extension" == "webm" ] && [ "$inputAudioCodec" != "" ] && [ "$inputAudioCodec" != "vorbis" ] && [ "$inputAudioCodec" != "opus" ]; then
		echo "error: unsupported audio codec ($inputAudioCodec) for webm container"
		exit 8
	fi
fi

fpsRatio=$(ffprobe -select_streams v:0 -show_entries stream=width,height,r_frame_rate "$inputVideo" 2>/dev/null | grep r_frame_rate= | cut -f2 -d=)
if [ -z "$fpsRatio" ]; then
	fpsOpt=""
else
	let fps=100*$fpsRatio; fps=$(echo $fps | sed 's/\(..\)$/.\1/')
	fpsOpt="-r $fps"
fi

mkdir "inputFrames" "outputFrames"
ffmpeg -i "$inputVideo" -qscale:v 1 -qmin 1 -qmax 1 -vsync 0 $inputFrames/frame%08d.jpg
realesrgan-ncnn-vulkan -f jpg -n $model -i "$inputFrames/" -o "$outputFrames/"
rm -rf $inputFrames
mogrify -resize $geom $outputFrames/*.jpg
ffmpeg $fpsOpt -i $outputFrames/frame%08d.jpg -i "$inputVideo" -map 0:v:0 -map 1:a:0 -c:v $vcodec -c:a copy "$outputVideo" #-pix_fmt yuv420p -vb 1M