import csv
import hashlib
import json
import os
import pathlib
import av
from sampling import make_plan

root=pathlib.Path('/eval'); clip=os.environ['CLIP']
spec=json.loads(pathlib.Path('/specs.json').read_text())[clip]
digest=hashlib.sha256(pathlib.Path('/source.mp4').read_bytes()).hexdigest()
assert digest==spec['sha256']
assert not (root/'inputs/sampling_plan.json').exists(),'Completed plan exists; do not overwrite'
pts=[]; dimensions=set(); corrupt=[]
with av.open('/source.mp4') as source:
    stream=source.streams.video[0]; stream.codec_context.thread_count=2
    meta=dict(codec=stream.codec_context.name,width=stream.width,height=stream.height,
              average_rate=str(stream.average_rate),declared_frames=stream.frames)
    for number,frame in enumerate(source.decode(video=0),1):
        assert frame.pts is not None
        pts.append(float(frame.pts*frame.time_base));dimensions.add((frame.width,frame.height))
        if frame.is_corrupt:corrupt.append(number)
assert len(dimensions)==1 and not corrupt,(dimensions,corrupt)
if meta['declared_frames']:assert meta['declared_frames']==len(pts)
plan=make_plan(pts,spec['display_bounds']);plan.update(clip=clip,source_sha256=digest,
    decode_errors=0,corrupt_frames=corrupt,source_probe=meta,known_loop_boundaries=spec['known_loop_boundaries'])
with (root/'inputs/source_map.csv').open('w',newline='') as f:
    w=csv.writer(f);w.writerow(['frame_1based','elapsed_seconds']);w.writerows((i+1,p-pts[0]) for i,p in enumerate(pts))
(root/'inputs/sampling_plan.json').write_text(json.dumps(plan,indent=2))
(root/'reports/source_probe.json').write_text(json.dumps(meta,indent=2))
print(json.dumps(dict(clip=clip,source_frames=len(pts),pose_samples=len(plan['target_seconds']),
                     display_samples=len(plan['display_samples']),display_bounds=plan['display_bounds'],probe=meta)),flush=True)