#!/usr/bin/env python
## Copyright (c) 2021, Alliance for Open Media. All rights reserved
##
## This source code is subject to the terms of the BSD 3-Clause Clear License and the
## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was
## not distributed with this source code in the LICENSE file, you can obtain it
## at aomedia.org/license/software-license/bsd-3-c-c/.  If the Alliance for Open Media Patent
## License 1.0 was not distributed with this source code in the PATENTS file, you
## can obtain it at aomedia.org/license/patent-license/.
##
__author__ = "maggie.sun@intel.com, ryanlei@meta.com"

import fileinput
import logging
import math
import os
from shutil import copyfile

import Utils
from AV2CTCVideo import AS_Downscaled_Clips
from Config import (
    AOMScaler,
    ContentPath,
    FFMPEG,
    HDRConvert,
    HDRToolsConfigFileTemplate,
    LoggerName,
    Platform,
)
from Utils import ExecuteCmd, GetShortContentName, md5

subloggername = "VideoScaler"
loggername = LoggerName + "." + "%s" % subloggername
logger = logging.getLogger(loggername)


# use AOMScaler to do image rescaling
def RescaleWithAom(clip, outw, outh, algo, outfile, num, LogCmdOnly):
    out_s = outw / math.gcd(outw, clip.width)
    in_s = clip.width / math.gcd(outw, clip.width)
    assert out_s == int(out_s) and in_s == int(in_s)

    scaling_str = "%d:%d:5" % (int(out_s), int(in_s))

    args = " -ieb:6 %s %d %s:c,d, %s %s" % (
        clip.file_path,
        num,
        scaling_str,
        scaling_str,
        outfile,
    )
    if outw > clip.width and outh > clip.height:
        args += " %dx%d" % (outw, outh)

    cmd = AOMScaler + args
    ExecuteCmd(cmd, LogCmdOnly)


# use ffmpeg to do image rescaling
def RescaleWithFfmpeg(clip, outw, outh, algo, outfile, num, LogCmdOnly):
    if clip.fmt == "420" and clip.bit_depth == 8:
        pix_fmt = "yuv420p"
    elif clip.fmt == "420" and clip.bit_depth == 10:
        pix_fmt = "yuv420p10le"
    else:
        print("Unsupported color format")

    args = (
        " -y -i %s -vf scale=%d:%d -pix_fmt %s -strict -1"
        " -sws_flags %s+accurate_rnd+full_chroma_int -sws_dither none"
        % (clip.file_path, outw, outh, pix_fmt, algo)
    )
    if algo == "lanczos":
        args += " -param0 5 "
    args += " -frames %d %s" % (num, outfile)
    cmd = FFMPEG + args
    ExecuteCmd(cmd, LogCmdOnly)


def GenerateCfgFile(clip, outw, outh, algo, outfile, num, configpath):
    contentBaseName = GetShortContentName(clip.file_name, False)
    cfg_filename = contentBaseName + ("_Scaled_%s_%dx%d.cfg" % (algo, outw, outh))
    fmt = 1
    if clip.fmt == "400":
        fmt = 0
    elif clip.fmt == "420":
        fmt = 1
    elif clip.fmt == "422":
        fmt = 2
    elif clip.fmt == "444":
        fmt = 3

    fps = 0
    if clip.fps_num == 0:
        fps = 0
    else:
        fps = (float)(clip.fps_num / clip.fps_denom)

    cfgfile = os.path.join(configpath, cfg_filename)
    copyfile(HDRToolsConfigFileTemplate, cfgfile)
    fp = fileinput.input(cfgfile, inplace=1)
    for line in fp:
        if "SourceFile=" in line:
            line = 'SourceFile="%s"\n' % clip.file_path
        if "OutputFile=" in line:
            line = 'OutputFile="%s"\n' % outfile
        if "SourceWidth=" in line:
            line = "SourceWidth=%d\n" % clip.width
        if "SourceHeight=" in line:
            line = "SourceHeight=%d\n" % clip.height
        if "OutputWidth=" in line:
            line = "OutputWidth=%d\n" % outw
        if "OutputHeight=" in line:
            line = "OutputHeight=%d\n" % outh
        if "SourceRate=" in line:
            line = "SourceRate=%4.3f\n" % fps
        if "SourceChromaFormat=" in line:
            line = "SourceChromaFormat=%d\n" % fmt
        if "SourceBitDepthCmp0=" in line:
            line = "SourceBitDepthCmp0=%d\n" % clip.bit_depth
        if "SourceBitDepthCmp1=" in line:
            line = "SourceBitDepthCmp1=%d\n" % clip.bit_depth
        if "SourceBitDepthCmp2=" in line:
            line = "SourceBitDepthCmp2=%d\n" % clip.bit_depth
        if "OutputRate=" in line:
            line = "OutputRate=%4.3f\n" % fps
        if "OutputChromaFormat=" in line:
            line = "OutputChromaFormat=%d\n" % fmt
        if "OutputBitDepthCmp0=" in line:
            line = "OutputBitDepthCmp0=%d\n" % clip.bit_depth
        if "OutputBitDepthCmp1=" in line:
            line = "OutputBitDepthCmp1=%d\n" % clip.bit_depth
        if "OutputBitDepthCmp2=" in line:
            line = "OutputBitDepthCmp2=%d\n" % clip.bit_depth
        if "NumberOfFrames=" in line:
            line = "NumberOfFrames=%d\n" % num
        print(line, end="")
    fp.close()
    return cfgfile


def RescaleWithHDRTool(
    clip, outw, outh, algo, outfile, num, cfg_path, LogCmdOnly=False
):
    cfg_file = GenerateCfgFile(clip, outw, outh, algo, outfile, num, cfg_path)
    args = " -f %s" % cfg_file
    cmd = HDRConvert + args
    ExecuteCmd(cmd, LogCmdOnly)


def VideoRescaling(
    method, clip, num, outw, outh, outfile, algo, cfg_path, LogCmdOnly=False
):
    if method == "hdrtool":
        RescaleWithHDRTool(clip, outw, outh, algo, outfile, num, cfg_path, LogCmdOnly)
    elif method == "aom":
        RescaleWithAom(clip, outw, outh, algo, outfile, num, LogCmdOnly)
    else:
        RescaleWithFfmpeg(clip, outw, outh, algo, outfile, num, LogCmdOnly)

    # add other tools for scaling here later


####################################################################################
##################### Major Functions ################################################
def GetDownScaledOutFile(
    clip, dnw, dnh, path, method, algo, ds_on_the_fly=True, ratio_idx=0
):
    contentBaseName = GetShortContentName(clip.file_name, False)
    dnscaledout = clip.file_path
    if clip.width != dnw or clip.height != dnh:
        if ds_on_the_fly:
            filename = contentBaseName + (
                "_Scaled_%s_%s_%dx%d.y4m" % (method, algo, dnw, dnh)
            )
            dnscaledout = os.path.join(path, filename)
        else:
            dnscaledout = (
                ContentPath
                + "/A1_downscaled/"
                + AS_Downscaled_Clips[contentBaseName][ratio_idx - 1]
            )

    return dnscaledout


def GetUpScaledOutFile(clip, outw, outh, method, algo, path):
    contentBaseName = GetShortContentName(clip.file_name, False)
    upscaledout = clip.file_path
    if clip.width != outw or clip.height != outh:
        filename = contentBaseName + (
            "_Scaled_%s_%s_%dx%d.y4m" % (method, algo, outw, outh)
        )
        upscaledout = os.path.join(path, filename)
    return upscaledout


def GetDownScaledMD5File(clip, dnw, dnh, path, method, algo):
    contentBaseName = GetShortContentName(clip.file_name, False)
    filename = contentBaseName + ".md5"
    if clip.width != dnw or clip.height != dnh:
        filename = contentBaseName + (
            "_Scaled_%s_%s_%dx%d.md5" % (method, algo, dnw, dnh)
        )
    dnscaledmd5 = os.path.join(path, filename)
    return dnscaledmd5


def CalculateDownScaledMD5(clip, dnw, dnh, path, method, algo, LogCmdOnly):
    dnScaleMD5 = GetDownScaledMD5File(clip, dnw, dnh, path, method, algo)
    if LogCmdOnly == 1:
        if Platform == "Linux":
            cmd = "md5sum %s &> %s" % (clip.file_path, dnScaleMD5)
        ExecuteCmd(cmd, 1)
    else:
        f = open(dnScaleMD5, "wt")
        dnScaledOut = GetDownScaledOutFile(clip, dnw, dnh, path, method, algo)
        MD5 = md5(dnScaledOut)
        f.write(MD5)
        f.close()


def DownScaling(method, clip, num, outw, outh, path, cfg_path, algo, LogCmdOnly=False):
    dnScaledOut = GetDownScaledOutFile(clip, outw, outh, path, method, algo)

    Utils.CmdLogger.write("::Downscaling\n")
    if clip.width != outw or clip.height != outh:
        # call separate process to do the downscaling
        VideoRescaling(
            method, clip, num, outw, outh, dnScaledOut, algo, cfg_path, LogCmdOnly
        )

    CalculateDownScaledMD5(clip, outw, outh, path, method, algo, LogCmdOnly)

    return dnScaledOut


def UpScaling(method, clip, num, outw, outh, path, cfg_path, algo, LogCmdOnly=False):
    upScaleOut = GetUpScaledOutFile(clip, outw, outh, method, algo, path)
    Utils.CmdLogger.write("::Upscaling\n")
    if clip.width != outw or clip.height != outh:
        # call separate process to do the upscaling
        VideoRescaling(
            method, clip, num, outw, outh, upScaleOut, algo, cfg_path, LogCmdOnly
        )
    return upScaleOut
