#!/usr/bin/env python3
""" Convert video files in a folder recursively to mp4 format

NB:
If you run this script repeatedly on .mp4 files with the same preset (e.g. veryfast) and crf (e.g. 23.0),
the file size will be reduced a bit each time. - I think it is mainly crf.
The first time you run it, the file size reduction is the most significant.
"""
import os
import sys
from shutil import which
import argparse

# https://trac.ffmpeg.org/wiki/Encode/H.264
DEFAULT_CRF = 23.0
BETTER_QUALITY_CRF = (
    18.0  # bigger file size. 18 results in almost the same size as .mod file
)
DEFAULT_PRESET = "medium"  # veryfast and medium (default) are good options. A slower preset will provide better compression (compression is quality per filesize)
MP4_PRESET = "veryfast"  # for .mp4 files, use veryfast. The quality and file size are almost the same as medium but much faster.
NEW_FILE_EXT = ".mp4"  # If source is .mp4, it will be src_new.mp4

default_video_extensions = (
    ".mp4",
    ".avi",  # Windows Media Player cannot play converted .mp4 after conversion, but K-Lite Codec Pack / MPC-HC can play it.
    ".mkv",
    ".mov",
    ".mod",
    ".wmv",
    ".flv",
    ".mpeg",
    ".mpg",
    ".webm",
    ".3gp",
)


def register_signal_ctrl_c():
    """Register signal handler for ctrl+c"""
    import signal
    import sys

    def signal_handler(sig, frame):
        print(" You pressed Ctrl+C! Exiting...")
        sys.exit(1)

    signal.signal(signal.SIGINT, signal_handler)


def normal_path(path: str, resolve_symlink=False):
    """return a normalized path

    resolve_symlink: resolve symbolic link to the actual path

    calls os.path.x:
    expanduser: ~/a.txt => /home/user/a.txt
    expandvars: '$HOME/mydir' => /home/user/mydir
    normpath: /usr/local//../bin/ => /usr/bin
    realpath: a.txt => /home/user/project/a.txt (full path);  /usr/local/bin/python3 (symlink) => /usr/bin/python3.8
    abspath: a.txt => /home/user/project/a.txt (full path); /usr/local/bin/python3 (symlink) => /usr/local/bin/python3;
            /usr/local//../bin/ => /usr/bin

    Note: abspath does half realpath + normpath.
    """
    from os.path import expanduser, expandvars, normpath, realpath, abspath

    if resolve_symlink:
        return realpath(normpath(expandvars(expanduser(path))))
    else:
        return abspath(expandvars(expanduser(path)))


def bashx(cmd, x=True, e=False):
    """
    run system cmd like bash -x

    Print the cmd with + prefix before executing
    Don't capture the output or error

    Arguments
    ---------
    cmd: string - command to run
    x:  When True, print the command with prefix + like shell 'bash -x' before running it
    e:  When True, exit the python program when returncode is not 0, like shell 'bash -e'.

    return
    ------
    CompletedProcess object with only returncode.

    Shell environment variables
    ---------------------------
    You can set the following environment variables which have the same effect of but overwrite argument options.
    BASH_CMD_PRINT=True     same as x=True
    BASH_EXIT_ON_ERROR=True same as e=True

    Usage example:
    ret = shx('ls')
    print(ret.returncode)

    Warning of using shell=True: https://docs.python.org/3/library/subprocess.html#security-considerations
    """
    from subprocess import run  # CompletedProcess
    import sys
    import os

    if sys.version_info >= (
        3,
        5,
    ):  # version_info is actually a tuple, so compare with a tuple
        x_env = os.getenv("BASH_CMD_PRINT", None)
        # BASH_CMD_PRINT overwrites x argument
        if x_env is not None:
            if x_env.lower() == "true":
                x = True
            elif x_env.lower() == "false":
                x = False

        e_env = os.getenv("BASH_EXIT_ON_ERROR", None)
        # BASH_EXIT_ON_ERROR overwrites x argument
        if e_env is not None:
            if e_env.lower() == "true":
                e = True
            elif e_env.lower() == "false":
                e = False

        if x:
            print("+ %s" % cmd)

        ret = run(cmd, shell=True)

        if e and ret.returncode != 0:
            print(
                "[Error] Command returns error code %s. Exiting the program..."
                % ret.returncode
            )
            sys.exit(1)
        else:
            return ret
    else:
        raise Exception("Require python 3.5 or above.")


def parse_cli_args():
    parser = argparse.ArgumentParser(
        description="Convert video files in a directory recursively or a file to mp4 format. If src is mp4, its file size will be reduced thanks to optimization. Requires ffmpeg pre-installed."
    )

    parser.add_argument(
        "-r",
        "--remove-original",
        action="store_true",
        help="remove original video file after conversion, default is False",
    )
    parser.add_argument(
        "-e",
        "--extension",
        action="append",
        help=(
            "Video file extensions to convert, one or more parameters, e.g. two extensions, -e .mp4 -e .avi, "
            f"default: {default_video_extensions}"
        ),
    )
    parser.add_argument(
        "-p",
        "--path",
        type=str,
        default=".",
        help="directory to search for video files or a video file to convert",
    )
    return parser.parse_args()


def convert_to_mp4(src_fullpath, dst_fullpath, remove_original=False, crf=23.0):
    """Convert a video file to mp4 format using ffmpeg

    return: return code of the ffmpeg command, 0 for success, non-zero for failure
    """
    import shutil

    _, ext = os.path.splitext(src_fullpath)

    if ext == ".mp4":
        cmd = f'ffmpeg -y -i "{src_fullpath}" -c:v libx264  -crf {crf} -preset {MP4_PRESET} "{dst_fullpath}"'
    else:
        cmd = f'ffmpeg -y -i "{src_fullpath}" -c:v libx264  -crf {crf} -preset {DEFAULT_PRESET} "{dst_fullpath}"'
    print("Converting %s to %s" % (src_fullpath, dst_fullpath))
    ret = bashx(cmd)
    if ret.returncode == 0:
        print("Converted %s to %s successfully.\n" % (src_fullpath, dst_fullpath))
        if remove_original is True:
            os.remove(src_fullpath)
            if ext.lower() == ".mp4":
                shutil.move(dst_fullpath, src_fullpath)
                print(
                    f"Removed original file {src_fullpath} and renamed {dst_fullpath} to {src_fullpath}"
                )
    else:
        print("Error: Failed to convert %s to %s" % (src_fullpath, dst_fullpath))
    return ret.returncode


def get_dst_fullpath(src_fullpath):
    f_noext, f_ext = os.path.splitext(src_fullpath)
    if f_ext.lower() in [".mp4"]:
        return f_noext + "_new" + NEW_FILE_EXT
    else:
        return f_noext + NEW_FILE_EXT


def error_exit(message: str):
    print(f"{message}")
    input("Press Enter to exit... Or close this window.")
    sys.exit(1)


def main():
    register_signal_ctrl_c()
    args = parse_cli_args()
    if which("ffmpeg") is None:
        print("Error: ffmpeg is not installed!")
        sys.exit()

    if args.extension:
        video_extensions = [str(e).lower() for e in args.extension]
    else:
        video_extensions = default_video_extensions

    print("video_extensions: ", video_extensions)
    target_path = normal_path(args.path)

    err_files = []
    ok_files = []
    found_video = False

    # if path is a file, convert it
    if os.path.isfile(target_path):
        src_fullpath = target_path
        f_noext, f_ext = os.path.splitext(src_fullpath)
        if f_ext.lower() in video_extensions:
            found_video = True
            dst_fullpath = get_dst_fullpath(src_fullpath)
            # set lower crf for better quality for some files. crf range 0 - 51
            # if f_ext.lower() in ('.mod'):
            #     crf = BETTER_QUALITY_CRF
            # else:
            crf = DEFAULT_CRF
            ret = convert_to_mp4(
                src_fullpath, dst_fullpath, args.remove_original, crf=crf
            )
            if ret == 0:
                ok_files.append(target_path)
            else:
                err_files.append(target_path)
        else:
            error_exit(
                f"Error: {args.path} is not a video file with supported extension."
            )
    # if path is a directory, convert all video files in it
    elif os.path.isdir(target_path):
        # # double check if user wants to remove original file after conversion
        # if not args.remove_original:
        #     answer = input(
        #         f"Do you want to remove original file after conversion? (y/n): "
        #     )
        #     if answer.lower() == "y":
        #         args.remove_original = True

        # if '.mod' in video_extensions:
        #     answer = input(f"Converting .mod to .mp4 may result in a worse quality, do you want to continue? (y/n): ")
        #     if answer.lower() != 'y':
        #         sys.exit()

        for root, dirs, files in os.walk(target_path):
            for f in files:
                f_noext, f_ext = os.path.splitext(f)
                if f_ext.lower() in video_extensions:
                    found_video = True
                    src_fullpath = os.path.join(root, f)
                    dst_fullpath = get_dst_fullpath(src_fullpath)
                    # set lower crf for better quality for some files. crf range 0 - 51
                    # if f_ext.lower() in ('.mod'):
                    #     crf = BETTER_QUALITY_CRF
                    # else:
                    crf = DEFAULT_CRF
                    ret = convert_to_mp4(
                        src_fullpath, dst_fullpath, args.remove_original, crf=crf
                    )
                    if ret == 0:
                        ok_files.append(src_fullpath)
                    else:
                        err_files.append(src_fullpath)
    else:
        error_exit(f"Error: {args.path} is not a valid file or directory.")

    print()
    if found_video is False:
        print(f"Don't find any video files in {target_path} to convert.")
    elif len(err_files) != 0:
        print("Below files cannot be converted:")
        for f in err_files:
            print(f)
    elif len(ok_files) != 0:
        print("搞定！Below files are converted to mp4 or compressed:")
        for f in ok_files:
            print(f)

    input("Press Enter to exit... Or close this window")


if __name__ == "__main__":
    main()
