#!/usr/bin/env python3
"""Render a seamless video of one of the pages: frames first, then ffmpeg.

    render_wallpaper.py <page.html> <out.mp4> <width> <height> [seconds] [fps]

Frames land in <out.mp4>.frames/ as PNG and are kept so a re-run resumes.
The page's loop is exactly 30 s, so keep the duration a multiple of that to
make the video loop without a seam."""
import os
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

page, out = Path(sys.argv[1]).resolve(), Path(sys.argv[2]).resolve()
width, height = int(sys.argv[3]), int(sys.argv[4])
seconds = float(sys.argv[5]) if len(sys.argv) > 5 else 30.0
fps = int(sys.argv[6]) if len(sys.argv) > 6 else 30
workers = min(4, os.cpu_count() or 2)  # 8 headless Firefox at once segfaults the GPU driver

frames = out.with_suffix('.frames')
frames.mkdir(parents=True, exist_ok=True)
profiles = [Path(f'/tmp/ff_wall_{i}') for i in range(workers)]
for profile in profiles:
    profile.mkdir(exist_ok=True)

total = int(round(seconds * fps))


def render(index):
    target = frames / f'{index:04d}.png'
    if target.exists() and target.stat().st_size > 20000:
        return 0.0
    url = f'{page.as_uri()}#shot&t={index / fps:.6f}&info=0'
    profile = profiles[index % workers]
    for attempt in range(4):
        started = time.time()
        result = subprocess.run(['firefox', '--headless', '--new-instance',
                                 '--profile', str(profile),
                                 f'--window-size={width},{height}',
                                 f'--screenshot={target}', url],
                                timeout=300,
                                stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if result.returncode == 0 and target.exists() and target.stat().st_size > 20000:
            return time.time() - started
        target.unlink(missing_ok=True)
        time.sleep(1 + attempt)
    raise RuntimeError(f'frame {index} failed')


began = time.time()
with ThreadPoolExecutor(workers) as pool:
    for done, _ in enumerate(pool.map(render, range(total)), 1):
        if done % 30 == 0 or done == total:
            elapsed = time.time() - began
            print(f'{done}/{total} frames, {elapsed:.0f}s elapsed, '
                  f'~{elapsed / done * (total - done):.0f}s left', flush=True)

out.parent.mkdir(parents=True, exist_ok=True)
subprocess.run(['ffmpeg', '-y', '-hide_banner', '-loglevel', 'warning',
                '-framerate', str(fps), '-i', str(frames / '%04d.png'),
                # Baseline on purpose: Fedora's Qt Multimedia decodes H.264 through
                # openh264, which rejects High profile ("DecodeFrame failed" in the
                # journal) and then Plasma shows a frozen frame.
                '-c:v', 'h264_nvenc', '-profile:v', 'baseline', '-preset', 'p7',
                '-rc', 'constqp', '-qp', '18', '-pix_fmt', 'yuv420p',
                '-movflags', '+faststart', str(out)], check=True)
print(f'{out} written ({out.stat().st_size / 1e6:.1f} MB)')