"""Render a film/episode, export subtitles, and create a chapter contact sheet."""
import argparse
import hashlib
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys

import av
import imageio_ffmpeg
from PIL import Image, ImageDraw
from tex_environment import prepare_tex

HERE = Path(__file__).resolve().parent
parser = argparse.ArgumentParser()
parser.add_argument('--quality', choices=['preview', 'final'], default='preview')
parser.add_argument('--scene', choices=['SlimeFilm', 'ParallelEpisode', 'RandomEpisode'], default='SlimeFilm')
args = parser.parse_args()
os.chdir(HERE)
output = HERE / 'output'
output.mkdir(exist_ok=True)
tools = output / '.tools'
tools.mkdir(exist_ok=True)
ffmpeg = tools / 'ffmpeg'
if not ffmpeg.exists():
    ffmpeg.symlink_to(imageio_ffmpeg.get_ffmpeg_exe())
env = dict(os.environ, PATH=str(tools)+os.pathsep+os.environ['PATH'])
env = prepare_tex(env)
quality = '-ql' if args.quality == 'preview' else '-qh'
fps = '15' if args.quality == 'preview' else '30'
subprocess.run([sys.executable, '-m', 'manim', quality, '--fps', fps,
                '--format', 'mp4', '-o', args.scene, 'scenes.py', args.scene],
               env=env, check=True)
source = next(p for p in (HERE / 'media/videos/scenes').rglob(f'{args.scene}.mp4')
              if ('480p15' if args.quality == 'preview' else '1080p30') in str(p))
# Final resolution is 1080p; use explicit 30 fps to keep authoring/rendering practical.
target = output / f'{args.scene}-{args.quality}.mp4'
shutil.copy2(source, target)
timeline = json.loads((output / f'{args.scene}.timeline.json').read_text())
with av.open(str(target)) as video:
    stream = video.streams.video[0]
    report = {'file': target.name, 'width': stream.width, 'height': stream.height,
              'fps': float(stream.average_rate), 'duration': float(stream.duration*stream.time_base),
              'audio_streams': len(video.streams.audio), 'chapters': timeline['chapters'],
              'caption_count': len(timeline['cues'])}
    report['source_sha256'] = {p: hashlib.sha256((HERE / p).read_bytes()).hexdigest()
                               for p in ('scenes.py', 'render.py', 'tex_environment.py', 'check_lower35.py', 'manim.cfg', 'uv.lock', 'data/benchmarks.json')}
    samples = [chapter['start']+3 for chapter in timeline['chapters']]
    samples += [cue['start']+1 for cue in timeline['cues']
                if cue['text'].startswith(('A 组', 'B 组', 'C 组', '同一 4096'))]
    shots = []
    for t in samples:
        video.seek(int(t/stream.time_base), stream=stream, backward=True)
        for frame in video.decode(stream):
            if float(frame.pts*stream.time_base) >= t:
                im = frame.to_image()
                im.thumbnail((640, 360))
                shots.append(im)
                break
    sheet = Image.new('RGB', (1280, ((len(shots)+1)//2)*390), '#0B1220')
    draw = ImageDraw.Draw(sheet)
    for i, im in enumerate(shots):
        x, y = i%2*640, i//2*390
        sheet.paste(im, (x,y))
        draw.text((x+12,y+365), f'Frame {i+1} | {samples[i]:.1f}s', fill='white')
    sheet.save(output / f'{args.scene}-{args.quality}-contact.png')
(output / f'{args.scene}-{args.quality}.render.json').write_text(json.dumps(report, ensure_ascii=False, indent=2)+'\n')
# Keep the subtitle/script matching this exact quality's frame-aligned timeline.
for ext in ('srt', 'timeline.json', 'narration.md'):
    shutil.copy2(output / f'{args.scene}.{ext}', output / f'{args.scene}-{args.quality}.{ext}')
print(json.dumps(report, ensure_ascii=False, indent=2))