#!/usr/bin/env python3
"""Décode une séquence de frames JSON depuis CKLive et encode en WebM via ffmpeg.

Usage: python encode_webm.py <frames_dir_or_pattern> <output.webm> [--fps 25] [--quality 32]
"""
import base64
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

import imageio_ffmpeg

FFMPEG = imageio_ffmpeg.get_ffmpeg_exe()


def decode_batch_files(json_files: list[Path], png_dir: Path) -> int:
    """Décode tous les JSON de frames vers PNGs numérotés."""
    png_dir.mkdir(parents=True, exist_ok=True)
    n = 0
    for jf in sorted(json_files):
        with open(jf, 'r', encoding='utf-8') as f:
            data = json.load(f)
        # data is a list of dataURLs OR a dict with 'frames'
        if isinstance(data, list):
            frames = data
        elif isinstance(data, dict):
            frames = data.get('frames') or list(data.values())
        else:
            continue
        for url in frames:
            if not isinstance(url, str) or not url.startswith('data:image'):
                continue
            png = base64.b64decode(url.split(',', 1)[1])
            with open(png_dir / f'frame_{n:04d}.png', 'wb') as f:
                f.write(png)
            n += 1
    return n


def encode_webm(png_dir: Path, output: Path, fps: int = 25, quality: int = 32) -> None:
    """Encode PNG sequence en WebM (VP9, loop friendly)."""
    cmd = [
        FFMPEG, '-y',
        '-framerate', str(fps),
        '-i', str(png_dir / 'frame_%04d.png'),
        '-c:v', 'libvpx-vp9',
        '-crf', str(quality),
        '-b:v', '0',
        '-pix_fmt', 'yuv420p',
        '-row-mt', '1',
        '-deadline', 'good',
        '-cpu-used', '2',
        '-an',
        str(output)
    ]
    subprocess.run(cmd, check=True)
    print(f'  encoded -> {output} ({output.stat().st_size//1024} KB)')


def main():
    if len(sys.argv) < 3:
        print('Usage: encode_webm.py <batches_glob_or_json> <output.webm> [--fps N] [--quality Q]')
        sys.exit(2)
    src = sys.argv[1]
    output = Path(sys.argv[2]).resolve()
    fps = 25
    quality = 32
    i = 3
    while i < len(sys.argv):
        if sys.argv[i] == '--fps':
            fps = int(sys.argv[i+1]); i += 2
        elif sys.argv[i] == '--quality':
            quality = int(sys.argv[i+1]); i += 2
        else:
            i += 1
    # Find json files
    if '*' in src or '?' in src:
        import glob
        files = [Path(p) for p in glob.glob(src)]
    elif Path(src).is_dir():
        files = sorted(Path(src).glob('*.json'))
    else:
        files = [Path(src)]
    if not files:
        print(f'!!! No JSON files matching {src}'); sys.exit(1)
    print(f'  decoding {len(files)} JSON files...')
    tmp = output.parent / f'_pngs_{output.stem}'
    if tmp.exists():
        shutil.rmtree(tmp)
    n = decode_batch_files(files, tmp)
    print(f'  decoded {n} frames')
    output.parent.mkdir(parents=True, exist_ok=True)
    encode_webm(tmp, output, fps=fps, quality=quality)
    # Cleanup PNGs
    shutil.rmtree(tmp)


if __name__ == '__main__':
    main()