#!/usr/bin/env python3
"""Select one decoded frame by an explicit time/index policy. Local files only."""
import argparse, hashlib, json, os, re, shutil, struct, subprocess, sys
from fractions import Fraction
from pathlib import Path

def filehash(path):
    value=hashlib.sha256()
    with path.open("rb") as source:
        for part in iter(lambda:source.read(1048576), b""):value.update(part)
    return value.hexdigest()

def invoke(ffmpeg, args):
    return subprocess.run([ffmpeg, '-hide_banner', '-nostdin', '-threads', '1', *args], check=True, capture_output=True, text=True, timeout=90)

def executable():
    value=os.environ.get('FFMPEG_BIN') or shutil.which('ffmpeg')
    if not value: raise ValueError('Set FFMPEG_BIN to an installed FFmpeg executable (tested: 6.0).')
    return value

def inventory(ffmpeg, source):
    if not source.is_file() or source.is_symlink(): raise ValueError('Input must be a regular local file, not a URL or symlink.')
    run=invoke(ffmpeg,['-copyts','-noautorotate','-i',str(source),'-map','0:v:0','-an','-sn','-dn','-vf','showinfo','-fps_mode','passthrough','-f','null','-'])
    match=re.search(r'config in time_base:\s*(\d+)/(\d+)',run.stderr)
    if not match: raise ValueError('No decoder/filter time base found; do not infer it from nominal FPS.')
    tb=Fraction(int(match[1]),int(match[2])); frames=[]
    for item in re.finditer(r'\bn:\s*(\d+)\s+pts:\s*(-?\d+)\s+pts_time:',run.stderr):
        n,pts=int(item[1]),int(item[2]); frames.append({'index':n,'pts':pts,'timestamp':Fraction(pts)*tb})
    if not frames or len(frames)>100000: raise ValueError('No frames or unsupported inventory size (>100000 frames).')
    if [f['index'] for f in frames]!=list(range(len(frames))): raise ValueError('Non-contiguous decoded indices.')
    if any(b['timestamp']<=a['timestamp'] for a,b in zip(frames,frames[1:])): raise ValueError('Missing, repeated, or non-monotonic PTS require separate repair/interpretation.')
    first=frames[0]['timestamp']
    for f in frames: f['relative']=f['timestamp']-first
    return tb,frames,run.stderr

def choose(frames,rule,time=None,index=None):
    if rule=='frame':
        if index is None or index<0 or index>=len(frames): raise ValueError('Zero-based frame index is outside decoded inventory.')
        return frames[index]
    if time is None: raise ValueError('Time policy needs --time in seconds; accepts a decimal or rational such as 1/5.')
    target=Fraction(time)
    if target<0 or target>frames[-1]['relative']: raise ValueError('Time outside first-to-last PTS range; final display duration is not inferred.')
    candidates=[f for f in frames if f['relative']<=target] if rule=='at' else [f for f in frames if f['relative']>=target]
    return candidates[-1] if rule=='at' else candidates[0]

def row(f,tb):
    return {'frameIndex':f['index'],'pts':f['pts'],'timeBase':str(tb),'timestampSecondsRational':str(f['timestamp']), 'timestampSeconds':float(f['timestamp']),'relativeSecondsRational':str(f['relative']),'relativeSeconds':float(f['relative'])}

def extract(ffmpeg,source,out,rule,time=None,index=None,orientation='display'):
    if out.suffix.lower()!='.png': raise ValueError('Output must use .png.')
    out.parent.mkdir(parents=True,exist_ok=True)
    receipt=out.with_suffix('.selection.json'); framehash=out.with_suffix('.framehash.txt'); probe=out.with_suffix('.inventory.json')
    if any(p.exists() or p.is_symlink() for p in [out,receipt,framehash,probe]): raise ValueError('Refusing to overwrite output or sidecars; select a fresh output name.')
    tb,frames,_=inventory(ffmpeg,source); selected=choose(frames,rule,time,index)
    filter_=f"select=eq(n\\,{selected['index']}),format=rgb24"
    orient=[] if orientation=='display' else ['-noautorotate']
    common=['-copyts',*orient,'-i',str(source),'-map','0:v:0','-an','-sn','-dn','-vf',filter_,'-fps_mode','passthrough','-frames:v','1']
    invoke(ffmpeg,[*common,'-c:v','png','-threads','1','-f','image2','-update','1','-n',str(out)])
    hashed=invoke(ffmpeg,[*common,'-c:v','rawvideo','-enc_time_base',f'{tb.numerator}:{tb.denominator}','-f','framehash','-hash','sha256','-'])
    data=[line for line in hashed.stdout.splitlines() if line and not line.startswith('#')]
    if len(data)!=1: raise ValueError('Selected framehash did not produce exactly one frame.')
    pixelhash=data[0].split(',')[-1].strip()
    pnghash=invoke(ffmpeg,['-i',str(out),'-map','0:v:0','-vf','format=rgb24','-frames:v','1','-c:v','rawvideo','-f','framehash','-hash','sha256','-'])
    pngdata=[line for line in pnghash.stdout.splitlines() if line and not line.startswith('#')]
    if len(pngdata)!=1 or pngdata[0].split(',')[-1].strip()!=pixelhash: raise ValueError('PNG decoded RGB hash does not match selected source frame.')
    result={'schemaVersion':1,'policy':rule,'requestedTime':time,'requestedIndex':index,'orientationPolicy':orientation,'selection':row(selected,tb),'inputSha256':filehash(source),'outputPngSha256':filehash(out),'decodedRgbSha256':pixelhash,'pixelMatchVerified':True,'outputDimensions':list(struct.unpack('>II',out.read_bytes()[16:24])),'timeOrigin':'relative to first decoded presentation timestamp','lastFrameDurationInferred':False}
    framehash.write_text(hashed.stdout);probe.write_text(json.dumps([row(f,tb) for f in frames],indent=2)+'\n');receipt.write_text(json.dumps(result,indent=2)+'\n')
    return result

def main():
    p=argparse.ArgumentParser(description=__doc__)
    p.add_argument('--input',required=True,type=Path);p.add_argument('--out',required=True,type=Path)
    p.add_argument('--rule',required=True,choices=['at','after','frame']);p.add_argument('--time');p.add_argument('--index',type=int)
    p.add_argument('--orientation',choices=['display','raw'],default='display');a=p.parse_args()
    if a.rule=='frame' and a.time is not None: p.error('frame uses --index, not --time')
    if a.rule!='frame' and a.index is not None: p.error('time policies use --time, not --index')
    try: print(json.dumps(extract(executable(),a.input,a.out,a.rule,a.time,a.index,a.orientation),indent=2))
    except (ValueError,subprocess.CalledProcessError,subprocess.TimeoutExpired) as e: print(f'Extraction stopped: {e}',file=sys.stderr);return 2
    return 0
if __name__=='__main__': sys.exit(main())