"""Caption-only repair. Requires PyAV 14.2.0, numpy, Pillow. No runtime capture."""
import av, json, shutil, hashlib, itertools, sys
import numpy as np
from fractions import Fraction
from pathlib import Path
from PIL import Image,ImageDraw,ImageFont
P=Path(__file__).resolve().parent
SRC=P/'source'
FONT='/usr/share/fonts/truetype/liberation/LiberationMono-Regular.ttf'
rows=[r for r in json.loads((SRC/'captures-all.json').read_text()) if not r['id'].startswith('K3-1440')]
def sha(p): return hashlib.sha256(p.read_bytes()).hexdigest()
def plane(p): return np.frombuffer(p,dtype=np.uint8).reshape(p.height,p.line_size)[:,:p.width]
def run(r):
 name=r['id'];original=SRC/(name+'.mp4');raw=P/'originals'/original.name
 if not raw.exists():shutil.copy2(original,raw)
 expected=next(x['sha256'] for x in json.loads((SRC/'manifest.json').read_text()) if x['name']==raw.name)
 assert sha(raw)==expected
 if original.exists():assert sha(raw)==sha(original)
 trans=json.loads((P/'inspection'/(name+'-caption-transitions.json')).read_text())
 first,end=trans[-2]['frame'],trans[-1]['frame']
 inp=av.open(str(raw));s=inp.streams.video[0];w,h=s.width,s.height
 y0=1604 if w==780 else 1744 # strictly inside existing strip: actual tops 1602 / 1741
 caption='Cancel, then retry; buttons and\nterminal show current state.' if w==780 else 'Cancel, then retry; buttons and terminal show current state.'
 band=Image.new('L',(w,h-y0),17);d=ImageDraw.Draw(band);d.multiline_text((16,10),caption,font=ImageFont.truetype(FONT,24),fill=255,spacing=4)
 band=np.array(band)
 # Demux packet durations separately, preserving original irregular frame times and final hold.
 durations={p.pts:p.duration for p in av.open(str(raw)).demux(video=0) if p.pts is not None}
 durations[max(durations)]=s.duration-max(durations) # source track end, not unreliable packet duration
 outpath=P/(name+'-caption-fixed.mp4');out=av.open(str(outpath),'w');v=out.add_stream('libx264',rate=30)
 v.width=w;v.height=h;v.pix_fmt='yuv420p';v.time_base=s.time_base;v.codec_context.time_base=s.time_base;v.codec_context.color_range=2
 v.options={'qp':'0','preset':'ultrafast','bf':'0','threads':'2'}
 count=0;pts=[]
 def mux(packet):
  packet.duration=durations[packet.pts];out.mux(packet)
 for i,f in enumerate(inp.decode(video=0)):
  nf=av.VideoFrame(w,h,'yuv420p')
  for a,b in zip(f.planes,nf.planes):plane(b)[:]=plane(a)
  nf.color_range=2;nf.pts=f.pts;nf.time_base=f.time_base
  if first<=i<end:
   plane(nf.planes[0])[y0:]=band
   for q in nf.planes[1:]:plane(q)[y0//2:]=128
  for packet in v.encode(nf):mux(packet)
  pts.append(f.pts);count+=1
 for packet in v.encode():mux(packet)
 out.close();inp.close()
 # Exhaustive decoded planar equality outside the permitted rectangle, and full equality elsewhere.
 a=av.open(str(raw));b=av.open(str(outpath));assert a.streams.video[0].duration*a.streams.video[0].time_base==b.streams.video[0].duration*b.streams.video[0].time_base,'track end changed';checks=[];changed=0
 for i,pair in enumerate(itertools.zip_longest(a.decode(video=0),b.decode(video=0))):
  f,g=pair;assert f is not None and g is not None,'frame count'
  assert Fraction(f.pts)*f.time_base==Fraction(g.pts)*g.time_base,'PTS changed'
  for j,(fp,gp) in enumerate(zip(f.planes,g.planes)):
   x,z=plane(fp),plane(gp);limit=(y0 if j==0 else y0//2) if first<=i<end else fp.height
   assert np.array_equal(x[:limit],z[:limit]),f'outside caption mismatch frame {i} plane {j}'
   changed+=int(np.count_nonzero(x!=z))
  checks.append(float(f.time))
 result={'id':name,'original':raw.name,'corrected':outpath.name,'originalSha256':sha(raw),'correctedSha256':sha(outpath),'originalBytes':raw.stat().st_size,'correctedBytes':outpath.stat().st_size,'frames':count,'replacementFrameRangeHalfOpen':[first,end],'replacementMediaRangeHalfOpen':[checks[first],checks[end]],'rectangle':[0,y0,w,h-y0],'actualCaptionTop':1602 if w==780 else 1741,'caption':caption,'verification':'Every decoded frame: exact PTS and track-end equality; exact Y/U/V sample equality outside rectangle during repair and entire frame otherwise. Full range tagged, lossless H.264 qp=0. No frame insertion, removal, splice, speed change, crop, app or terminal masking.','changedSamplesInsideCaption':changed,'firstPts':pts[0],'lastPts':pts[-1],'timeBase':str(s.time_base)}
 (P/(name+'-verification.json')).write_text(json.dumps(result,indent=2));print(json.dumps(result),flush=True)
if __name__=='__main__':
 for r in rows:
  if len(sys.argv)==1 or r['id'] in sys.argv[1:]:run(r)