import argparse
import os
import pathlib
import shlex
import subprocess

import yaml

# Hannes' AV1 wrapper script
# Version 2.0.1

# This is a good guide for optimizing AV1 parameters with the AOM reference encoder:
# https://www.reddit.com/r/AV1/comments/t59j32/encoder_tuning_part_4_a_2nd_generation_guide_to/

parser = argparse.ArgumentParser(prog="hav1w")
parser.add_argument("input")
parser.add_argument("output")
parser.add_argument("preset")
parser.add_argument("--crf")
parser.add_argument("--ba")
parser.add_argument("--alayout")
parser.add_argument("--cpu")
parser.add_argument("--mf")
parser.add_argument("--dnl")
parser.add_argument("--arnr-strength")
parser.add_argument("--crop")
parser.add_argument("--no10bit", action="store_true")
parser.add_argument("--dry", action="store_true")
args = parser.parse_args()

preset_path = os.path.join(pathlib.Path(__file__).parent.resolve(), "av1_presets.yml")
with open(preset_path, "r") as stream:
    presets = yaml.safe_load(stream)
preset = presets[args.preset]

if args.crf is not None:
    preset["crf"] = args.crf
if args.ba is not None:
    preset["ba"] = args.ba
if args.alayout is not None:
    preset["alayout"] = args.alayout
if args.cpu is not None:
    preset["cpu"] = args.cpu
if args.mf is not None:
    preset["mf"] = args.mf
if args.dnl is not None:
    preset["dnl"] = args.dnl
if args.arnr_strength is not None:
    preset["arnr_strength"] = args.arnr_strength

cmd = ["ffmpeg", "-i", args.input, "-c:v", "libaom-av1", "-c:a", "libopus", "-c:s", "copy", "-map", "0"]

# AV1 Settings

# Constant rate factor, determines target quality
cmd += ["-crf", preset["crf"], "-b:v", "0"]

# Set minimum keyframe interval to 12 to prevent placing too many keyframes... just in case
cmd += ["-keyint_min", "12"]

# Set lookahead to 48 frames
# Reddit post says it's worth it...
cmd += ["-lag-in-frames", "48"]

# Temporal filtering parameters
# Set arnr_strength to 0 for animation and low variance stuff, 1 for content with more variance and fast motion scenes
cmd += ["-arnr-strength", preset["arnr_strength"], "-arnr-max-frames", "3"]

# Disable deltaq-mode for higher fidelity
aom_params = ["deltaq-mode=0"]

# Enable quantization matrices
# The Reddit post told me to do so, no matter what :P
aom_params += ["enable-qm=1"]

# Grain synthesis
# Set to -1 for disabling grain synthesis: makes sense for animation content (animes, ...) as
# a) this is not a strength of the aom encoder
# b) animation content usually does not have a lot of noise
#
# If you need to do grain synthesis on animation content, maybe SVT-AV1 is a better choice?
# And read this: https://www.reddit.com/r/AV1/comments/n4si96/encoder_tuning_part_3_av1_grain_synthesis_how_it/
# TEMPORARILY DISABLED DUE TO https://bugs.chromium.org/p/aomedia/issues/detail?id=2768
preset["dnl"] = "-1"
if preset["dnl"] != "-1":
    aom_params += ["enable-dnl-denoising=0", "denoise-noise-level=" + preset["dnl"]]

cmd += ["-aom-params", ":".join(aom_params)]

if not args.no10bit:  # Encode to 10 bit per default
    cmd += ["-pix_fmt", "yuv420p10le"]

if args.crop is not None:
    cmd += ["-vf", "crop=" + args.crop]

# Encoding efficiency settings (tiles also affect decoding performance)
cmd += ["-cpu-used", preset["cpu"], "-row-mt", "1", "-tiles", "2x2"]

# Audio/Opus settings
cmd += ["-b:a", preset["ba"], "-af", f"aformat=channel_layouts={preset['alayout']}", "-mapping_family", preset["mf"]]

cmd += [args.output]

if args.dry:
    print(" ".join(list(map(lambda x:shlex.quote(x), cmd))))
    exit(0)

p = subprocess.run(cmd)
exit(p.returncode)