#!/usr/bin/env python3
"""Encode timestamped PNG captures and compare EVERY decoded frame to its source."""
import bisect,json,math,os,subprocess,sys
from pathlib import Path
import numpy as np
from PIL import Image,ImageDraw,ImageFont
run=Path(sys.argv[1]);fps=10
meta=json.loads((run/'metadata.json').read_text());frames=json.loads((run/'frames.json').read_text())
start=next(x for x in meta['timeline'] if x['action']=='recording-start')
end=next(x for x in meta['timeline'] if x['action']=='recording-end')
n=math.ceil(end['timeMs']/1000*fps);times=[f['timeMs'] for f in frames]
assert 1000*(len(frames)-1)/(times[-1]-times[0])>=10,'PNG capture below required 10 fps'
inputs=run/'encode-input';inputs.mkdir(exist_ok=True)
mapping=[]
for i in range(n):
 t=i*1000/fps; k=bisect.bisect_left(times,t)
 k=min(range(max(0,k-1),min(len(times),k+1)),key=lambda k:abs(times[k]-t))
 dest=inputs/f'{i:06d}.png'
 if not dest.exists():os.link(run/frames[k]['file'],dest)
 mapping.append({'frame':i,'videoMs':t,'source':frames[k]['file'],'sourceMs':times[k]})
name=f"{meta['side']}-{meta['scenario']}";output=run/(name+'.mp4')
cmd=['ffmpeg','-y','-v','error','-framerate',str(fps),'-i',str(inputs/'%06d.png'),'-c:v','libx264','-crf','0','-preset','slow','-tune','stillimage','-pix_fmt','yuv444p','-movflags','+faststart',str(output)]
subprocess.run(cmd,check=True)
w,h=meta['viewport'].values();p=subprocess.Popen(['ffmpeg','-v','error','-i',str(output),'-f','rawvideo','-pix_fmt','rgb24','pipe:1'],stdout=subprocess.PIPE)
size=w*h*3;maximum=0;perframe=[];cache={};rawdir=run/'verified-decoded';rawdir.mkdir(exist_ok=True)
for i,m in enumerate(mapping):
 b=bytearray()
 while len(b)<size:
  chunk=p.stdout.read(size-len(b))
  if not chunk:break
  b.extend(chunk)
 if len(b)!=size:raise RuntimeError(f'decoded short frame {i}')
 decoded=np.frombuffer(b,dtype=np.uint8).reshape(h,w,3)
 if m['source'] not in cache:
  cache={m['source']:np.array(Image.open(run/m['source']).convert('RGB'),dtype=np.int16)}
 delta=np.abs(decoded.astype(np.int16)-cache[m['source']]);mx=int(delta.max());maximum=max(maximum,mx)
 perframe.append({'frame':i,'videoMs':m['videoMs'],'maxRgbDelta':mx,'meanRgbDelta':float(delta.mean())})
 Image.fromarray(decoded).save(rawdir/f'{i:06d}.png')
assert not p.stdout.read(1),'Unexpected extra decoded frame'
assert p.wait()==0
assert maximum<=4,f'Unexpected encoder delta {maximum}'
steps=meta['timeline'][meta['timeline'].index(start):-1]
font=ImageFont.truetype('/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf',16)
thumbs=[]
for step in steps:
 t=max(0,step['timeMs']+ (300 if step['action'].startswith('request ') else -150 if step['action'].startswith('still ') else 100));i=min(n-1,round(t/1000*fps))
 im=Image.open(rawdir/f'{i:06d}.png').convert('RGB');b=step['panelBounds']
 if b:
  x=max(0,math.floor(b['x'])-8);y=max(0,math.floor(b['y'])-8)
  im=im.crop((x,y,min(w,math.ceil(b['x']+b['width'])+8),h))
 if im.width>470:im=im.resize((470,round(im.height*470/im.width)))
 label=f"{i/fps:.1f}s | {step['action']}"
 thumbs.append((label,im))
cols=3;cw=max(490,max(im.width for _,im in thumbs)+20);ch=max(im.height for _,im in thumbs)+70
sheet=Image.new('RGB',(cols*cw,math.ceil(len(thumbs)/cols)*ch),'#252525');draw=ImageDraw.Draw(sheet)
for j,(label,im) in enumerate(thumbs):
 x=j%cols*cw;y=j//cols*ch;draw.text((x+8,y+6),label,fill='white',font=font);sheet.paste(im,(x+8,y+42))
sheet.save(run/(name+'-contact-sheet.png'))
report={'recording':name,'encoding':cmd,'width':w,'height':h,'fps':fps,'sourceFrames':len(frames),'durationSeconds':n/fps,'sourceAverageFps':1000*(len(frames)-1)/(times[-1]-times[0]),'sourceMaxGapMs':max(b-a for a,b in zip(times,times[1:])),'decodedFramesVerified':len(perframe),'maxRgbDelta':maximum,'narratedSteps':steps,'perFrame':perframe,'frameMapping':mapping}
(run/'verification.json').write_text(json.dumps(report,indent=2))
print(json.dumps({k:report[k] for k in ['recording','sourceFrames','sourceAverageFps','sourceMaxGapMs','decodedFramesVerified','durationSeconds','maxRgbDelta']}),flush=True)