"""Render all article companions with bundled fonts, plus matching WebM/posters.

Use --force after source/font changes; --output supports review before replacing
published assets. Each scene has its own cache so parallel renders cannot race.
"""
import argparse
from concurrent.futures import ThreadPoolExecutor
from fractions import Fraction
import hashlib
import os
import pathlib
import shutil
import subprocess
import sys
import av

ROOT = pathlib.Path(__file__).resolve().parents[1]
CLASSES = {'mesh':'LampMesh','points':'LampPoints','voxels':'LampVoxels',
           'splats':'LampGaussians','nerf':'LampNerf','motion':'LampMotion',
           'time-field':'LampTimeField','time-splats':'LampTimeSplats'}

def write_revision(out):
    # Refresh browser/CDN URLs whenever any of the delivered media bytes change.
    files=sorted(p for p in out.iterdir() if p.suffix in {'.mp4','.webm','.webp'})
    digest=hashlib.sha256(b''.join(p.name.encode()+p.read_bytes() for p in files)).hexdigest()[:16]
    (ROOT/'src/lib/media.ts').write_text(
        '// Generated by scripts/render-media.py after media exports.\n'
        + f'export const MEDIA_REVISION="{digest}";\n'
        + 'export const mediaUrl=(filename:string)=>`/media/${filename}?v=${MEDIA_REVISION}`;\n')

def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--force', action='store_true')
    parser.add_argument('--output', type=pathlib.Path, default=ROOT/'public/media')
    parser.add_argument('--jobs', type=int, choices=[1,2], default=2)
    parser.add_argument('--scenes', nargs='+', choices=CLASSES, default=list(CLASSES))
    args = parser.parse_args()
    out = args.output.resolve()
    out.mkdir(parents=True, exist_ok=True)
    work = ROOT/'media/typography-render'
    work.mkdir(parents=True, exist_ok=True)

    def render(slug):
        cls = CLASSES[slug]
        dest = out/f'{slug}.mp4'
        cache = work/slug
        cache.mkdir(exist_ok=True)
        if args.force or not dest.exists():
            print('RENDERING', slug, flush=True)
            env = dict(os.environ, XDG_CACHE_HOME=str(work/'font-cache'))
            with (cache/'render.log').open('w') as log:
                subprocess.run([sys.executable,'-m','manim','-qh','--fps','24',
                                '--disable_caching','--media_dir',str(cache),
                                str(ROOT/'animations/lamp_scenes.py'),cls],
                               check=True,cwd=ROOT,env=env,stdout=log,stderr=log)
            videos=list((cache/'videos/lamp_scenes').rglob(cls+'.mp4'))
            assert len(videos)==1, videos
            shutil.copy2(videos[0],dest)
        with av.open(str(dest)) as inp, av.open(str(out/f'{slug}.webm'),'w') as target:
            src=inp.streams.video[0]
            stream=target.add_stream('libvpx-vp9',rate=24)
            stream.width=src.width;stream.height=src.height;stream.pix_fmt='yuv420p'
            stream.options={'crf':'28','b':'0','deadline':'realtime','cpu-used':'6'}
            last=None
            for i,frame in enumerate(inp.decode(video=0)):
                last=frame;frame.pts=i;frame.time_base=Fraction(1,24)
                for packet in stream.encode(frame):target.mux(packet)
            for packet in stream.encode():target.mux(packet)
            if last is None:raise RuntimeError(f'Empty render: {dest}')
            last.to_image().save(out/f'{slug}-poster.webp',quality=94)
        print('EXPORTED',slug,flush=True)
        return slug

    with ThreadPoolExecutor(max_workers=args.jobs) as pool:
        list(pool.map(render,args.scenes))
    if out==(ROOT/'public/media').resolve():write_revision(out)
    print('Complete:',out,flush=True)

if __name__=='__main__':
    main()