#!/usr/bin/env python

__description__ = """
Compiles images into one .webm file while retaining original image's frame size & aspect ratio.
The resulting file will have dynamic frame size.

ffmpeg is required.

How it works:
    1. Asynchronously converts images to .webm files.
    2. Concatenates resulting .webm files into single .webm.

"""

__epilog__ = """
Heavily inspired by https://github.com/kvdomingo/webm-dr-py.
"""

from shutil import copy2
from subprocess import run, PIPE
from pathlib import Path
from shlex import split
from multiprocessing import Pool
from os import cpu_count


import tempfile
import argparse
import time


def is_image_fn(image: Path) -> bool:
    return image.stem.isdigit() and len(image.stem) == 8 and image.is_file()


def sanitize(args) -> tuple[Path]:
    if not args.input.is_dir():
        raise ValueError(f'Input path is not a directory or is not available.')

    infolder = args.input.resolve()

    if args.output:
        outfile = args.output
    else:
        outfile = args.input

    outfile = outfile.with_suffix(
        '.webm' if outfile.suffix == '.webm' else outfile.suffix + '.webm')

    outfile = outfile.resolve()

    if not args.force:
        if outfile.is_file():
            raise RuntimeError(
                f'Not overwriting already existing file "{outfile}".')

    return infolder, outfile


def fix_demuxer(demuxer: Path, tmpdir: Path) -> str:
    """Makes file references in demuxer file be relative to the temp folder."""

    demuxer_path = tmpdir / 'demuxer.txt'

    with open(demuxer_path, 'w') as temp_demuxer:
        with open(demuxer) as broken_demuxer:
            for line in broken_demuxer:
                if line.startswith("file "):
                    temp_demuxer.write(
                        f"file '{tmpdir/Path(split(line)[1])}'\n")
                else:
                    temp_demuxer.write(line)


def generate_demuxer(tmpdir: Path) -> str:

    demuxer_path = tmpdir / 'demuxer.txt'

    with open(demuxer_path, 'w') as demuxer:
        demuxer.writelines(f"file '{image}'\n" for image in tmpdir.iterdir()
                           if is_image_fn(image))


def frame2webm(arg):
    image, tmpdir, args = arg

    webm = (tmpdir / image.name).with_suffix('.webm')

    cmd = run(
        [
            'ffmpeg',
            '-y',
            '-hide_banner',
            '-loglevel', 'error',
            '-f', 'image2',
            '-framerate', args.framerate,
            '-i', image.resolve(),
            '-c:v', args.encoder,
            '-crf', str(args.crf),
            '-cpu-used', str(args.cpu_used),
            '-pix_fmt', args.pix_fmt,
            webm
        ],
        text=True, stderr=PIPE
    )

    return cmd


def frames2webms(infolder: Path, tmpdir: Path, args):

    convert_args = ((img, tmpdir, args)
                    for img in infolder.iterdir() if is_image_fn(img))

    start_time = time.time_ns()

    with Pool(processes=args.job_count) as pool:
        for cmd in pool.imap_unordered(frame2webm, convert_args):
            if cmd.returncode != 0:
                raise RuntimeError('ffmpeg: ' + cmd.stderr)

    print(f"Took: {(time.time_ns() - start_time) * 1e-9}s.")


def webms2webm(outfile: Path, tmpdir: Path):
    cmd = run(
        [
            'ffmpeg',
            '-y',
            '-hide_banner',
            '-loglevel', 'error',
            '-f', 'concat',
            '-safe', '0',
            '-i', tmpdir / "demuxer.txt",
            '-c:v', 'copy',
            outfile
        ],
        text=True, stderr=PIPE
    )

    if cmd.returncode != 0:
        raise RuntimeError('ffmpeg: ' + cmd.stderr)


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=__description__,
        epilog=__epilog__,
        formatter_class=argparse.RawDescriptionHelpFormatter)

    parser.add_argument("-o", "--output", type=Path,
                        help="Path to save output file to.")
    parser.add_argument("-f", "--force", action="store_true",
                        default=False, help="Overwrite output file.")
    parser.add_argument("-d", "--demuxer", type=Path,
                        help="Path to custom concat demuxer file. Paths are relative to temporary folder.")
    parser.add_argument("-r", "--framerate", type=str,
                        default='1')
    parser.add_argument("-e", "--encoder", type=str,
                        default="libaom-av1", help="Encoder used to encode video. Tested with: libaom-av1 (default), libvpx-vp9.")
    parser.add_argument("--crf", type=float,
                        default=20, help="Quality level. Less is better quality but larger file. 0 is lossless.")
    parser.add_argument("--job-count", type=int,
                        default=cpu_count() // 4, help="Number of jobs converting individual images to webms.")
    parser.add_argument("--pix_fmt", type=str,
                        default='yuv420p')
    parser.add_argument("--cpu-used", type=int,
                        default=8)
    parser.add_argument("--export-demuxer", type=Path,
                        help="Export generated concat demuxer to specified file.")
    parser.add_argument("input", type=Path,
                        help="Path to folder containing <%%8d>.<ext> formatted images.")

    args = parser.parse_args()

    infolder, outfile = sanitize(args)

    with tempfile.TemporaryDirectory() as tmpdir:
        tmpdir = Path(tmpdir)

        print("Frames -> webms...")
        frames2webms(infolder, tmpdir, args)

        print("Demuxer...")
        if args.demuxer:
            fix_demuxer(args.demuxer, tmpdir)
        else:
            generate_demuxer(tmpdir)

        if args.export_demuxer:
            copy2(tmpdir / "demuxer.txt", args.export_demuxer)

        print("webms -> webm...")
        webms2webm(outfile, tmpdir)

    print(outfile)