#!/usr/bin/env python3
"""Create original CFR/VFR fixtures and execute timestamp-selection checks."""
import argparse, hashlib, importlib.util, json, os, platform, shutil, struct, subprocess, sys
from datetime import datetime,timezone
from pathlib import Path
from fractions import Fraction
sys.dont_write_bytecode=True
here=Path(__file__).resolve().parent
spec=importlib.util.spec_from_file_location('selector',here/'extract-frame.py');selector=importlib.util.module_from_spec(spec);spec.loader.exec_module(selector)

def run(work):
 ffmpeg=selector.executable()
 if work.exists(): raise ValueError('Output directory already exists; use a fresh --out-dir to preserve prior results.')
 work.mkdir(parents=True); commands=[];scenarios=[]
 script_hashes={name:hashlib.sha256((here/name).read_bytes()).hexdigest() for name in ['extract-frame.py','run-fixtures.py']}
 def media(args):
  commands.append(['ffmpeg',*[str(x).replace(str(work),'OUTPUT') for x in args]]);return selector.invoke(ffmpeg,args)
 def raw(n):
  # Unique coloured frame-index tiles and moving stripe; original deterministic RGB data, no third-party picture.
  data=bytearray()
  for i in range(n):
   for y in range(180):
    for x in range(320):
     color=(20+i*15,38+(x//32)*12,65+(y//30)*18)
     if 20+i*20<=x<34+i*20 and 25<=y<155:color=(235,173,86)
     if 228<=x<306 and 18<=y<68:
      bit=(x-228)//13
      color=(208,235,225) if ((i>>bit)&1) else (34,62,78)
     data.extend(color)
  return bytes(data)
 for name,pts,rate in [('cfr',list(range(0,1000,100)),10),('vfr',[0,40,120,280,320,600,640,920],25)]:
  source=work/(name+'.rgb');source.write_bytes(raw(len(pts)))
  args=['-f','rawvideo','-pixel_format','rgb24','-video_size','320x180','-framerate',str(rate),'-i',str(source),'-an']
  if name=='vfr':
   expr=str(pts[-1])
   for n in range(len(pts)-2,-1,-1):expr=f'if(eq(N,{n}),{pts[n]},{expr})'
   args+=['-vf','settb=1/1000,setpts='+expr.replace(',','\\,')]
  args+=['-c:v','ffv1','-level','3','-threads','1','-enc_time_base','1:1000','-fps_mode','passthrough','-n',str(work/(name+'.mkv'))]
  media(args);source.unlink()
  tb,frames,log=selector.inventory(ffmpeg,work/(name+'.mkv'));assert tb==Fraction(1,1000)
  assert [f['pts'] for f in frames]==pts
  (work/(name+'-inventory.json')).write_text(json.dumps([selector.row(f,tb) for f in frames],indent=2)+'\n')
  (work/(name+'-showinfo.txt')).write_text(log.replace(str(work),'OUTPUT'))
  hashed=media(['-copyts','-noautorotate','-i',str(work/(name+'.mkv')),'-map','0:v:0','-an','-vf','format=rgb24','-c:v','rawvideo','-enc_time_base','1:1000','-fps_mode','passthrough','-f','framehash','-hash','sha256','-'])
  (work/(name+'-all.framehash.txt')).write_text(hashed.stdout);scenarios.append(name+'-pts-inventory')
 cases=[('cfr','at','0.25',None,2,200),('cfr','after','0.25',None,3,300),('cfr','frame',None,3,3,300),('vfr','at','0.20',None,2,120),('vfr','after','0.20',None,3,280),('vfr','frame',None,3,3,280),('vfr','at','0.28',None,3,280),('vfr','after','0.28',None,3,280)]
 records=[]
 for i,(name,rule,time,index,n,pts) in enumerate(cases):
  out=work/f'{name}-{rule}-{i}.png';r=selector.extract(ffmpeg,work/(name+'.mkv'),out,rule,time,index)
  assert r['selection']['frameIndex']==n and r['selection']['pts']==pts and r['pixelMatchVerified'];records.append(r);scenarios.append(f'{name}-{rule}-{time if time is not None else index}')
 _,f,_=selector.inventory(ffmpeg,work/'vfr.mkv')
 for label,rule,time,index in [('negative-time','at','-0.1',None),('beyond-last-pts','after','1.1',None),('frame-out-of-range','frame',None,8)]:
  try:selector.choose(f,rule,time,index)
  except ValueError:scenarios.append(label)
  else:raise AssertionError(label)
 # A real rotation metadata case verifies extraction policy without implying every matrix is a simple 90-degree rotation.
 media(['-i',str(work/'cfr.mkv'),'-map','0:v:0','-an','-c:v','libx264','-crf','0','-pix_fmt','yuv444p','-threads','1','-n',str(work/'rotation-base.mp4')])
 media(['-i',str(work/'rotation-base.mp4'),'-map','0:v:0','-c','copy','-metadata:s:v:0','rotate=90','-n',str(work/'rotation-90.mp4')]);(work/'rotation-base.mp4').unlink()
 for orientation,dims in [('display',(180,320)),('raw',(320,180))]:
  out=work/f'rotation-{orientation}.png';r=selector.extract(ffmpeg,work/'rotation-90.mp4',out,'frame',index=0,orientation=orientation)
  assert struct.unpack('>II',out.read_bytes()[16:24])==dims
  records.append(r);scenarios.append('rotation-'+orientation)
 # Starting with a stream offset verifies that CLI times are relative to the first decoded PTS, not container absolute times.
 media(['-copyts','-itsoffset','2','-i',str(work/'vfr.mkv'),'-map','0:v:0','-c','copy','-n',str(work/'vfr-offset.mkv')])
 r=selector.extract(ffmpeg,work/'vfr-offset.mkv',work/'vfr-offset-after.png','after',time='0.20');assert r['selection']['pts']==2280 and r['selection']['relativeSecondsRational']=='7/25';records.append(r);scenarios.append('nonzero-start-relative-time')
 tool=media(['-version']).stdout
 (work/'tool-version.txt').write_text(tool)
 tables={}
 for name in ['cfr','vfr']:
  tables[name]=json.loads((work/(name+'-inventory.json')).read_text())
 with (work/'timestamp-table.csv').open('w') as out:
  out.write('fixture,frame_index,pts,time_base,timestamp_seconds,relative_seconds\n')
  for name,rows in tables.items():
   for f in rows:out.write(f"{name},{f['frameIndex']},{f['pts']},{f['timeBase']},{f['timestampSeconds']},{f['relativeSeconds']}\n")
 assert script_hashes=={name:hashlib.sha256((here/name).read_bytes()).hexdigest() for name in script_hashes}, 'Scripts changed during fixture execution'
 receipt={'schemaVersion':1,'scriptHashes':script_hashes,'recipe':'frame-extraction-timebase-validation','executedAt':datetime.now(timezone.utc).isoformat(),'pythonVersion':platform.python_version(),'platform':platform.system()+' '+platform.machine(),'ffmpegVersion':tool.splitlines()[0],'ffmpegBinarySha256':hashlib.sha256(Path(ffmpeg).read_bytes()).hexdigest(),'sourceType':'original synthetic RGB tiles; FFV1 CFR/VFR and H.264 rotation fixture','noModelGeneration':True,'productionTaskRun':False,'fixtureOnly':True,'passedScenarios':scenarios,'selectionRecords':records,'timestampTables':tables,'commands':commands,'checksums':{str(p.relative_to(work)):hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(work.iterdir()) if p.is_file()}}
 (work/'execution-receipt.json').write_text(json.dumps(receipt,indent=2)+'\n');print(json.dumps({'passed':len(scenarios),'scenarios':scenarios,'ffmpegVersion':receipt['ffmpegVersion'],'fixtureBytes':sum(p.stat().st_size for p in work.iterdir())},indent=2))
if __name__=='__main__':
 parser=argparse.ArgumentParser(description=__doc__);parser.add_argument('--out-dir',type=Path,default=here/'run-output');args=parser.parse_args()
 try:run(args.out_dir.resolve())
 except Exception as e:print(f'Fixture failed: {e}',file=sys.stderr);sys.exit(1)