import os,sys,json,csv,pathlib,itertools,hashlib,subprocess
from fractions import Fraction
from unittest.mock import patch
import av,numpy as np
from PIL import Image,ImageDraw,ImageFont
sys.path.insert(0,'/opt/contextra')
from app.fight_runtime import FightScenario,FightSettings
from app.fall_runtime import FallScenario,FallSettings
from app.contracts.frames import FrameEnvelope,PoseObservation
from app.pose_config import POSE_MODEL
from app.event_cooldown import EventCooldown
from app.fight_r5.geometry import bodies
from app.intrusion import inside

R=pathlib.Path('/work');C=pathlib.Path('/cache');O=R/'Contextra_20Clips_Engine_Review'
os.environ['CONTEXTRA_FIGHT_DISPLAY_POLICY']='/work/display-policy.json'
font='/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf'
F=ImageFont.truetype(font,36);H=ImageFont.truetype(font,48)
COLORS={'RED_CURRENT':(240,55,65),'YELLOW_HOLD':(245,195,40),'BLUE_CANDIDATE':(30,145,240)}
def plain(x):
 if hasattr(x,'tolist'):return x.tolist()
 if isinstance(x,set):return sorted(x)
 raise TypeError(type(x).__name__)
def sha(p):
 h=hashlib.sha256()
 with p.open('rb') as f:
  for b in iter(lambda:f.read(1048576),b''):h.update(b)
 return h.hexdigest()
def save(p,d):p.write_text(json.dumps(d,ensure_ascii=False,indent=2,default=plain))
def one(key,kind):
 out=O/kind/key;out.mkdir(parents=True,exist_ok=True)
 verify=out/'validation.json'
 if verify.exists():
  v=json.loads(verify.read_text());assert sha(out/v['video'])==v['sha256'];return v
 rows=json.loads((C/'clips'/key/'inputs/pose.json').read_text())
 with (C/'metadata/observation_maps'/f'{key}.csv').open(encoding='utf-8-sig') as f:maps=list(csv.DictReader(f))
 clock=[0.];ids=itertools.count(1)
 cfg=dict(enabled=True,engine='r5' if kind=='FIGHT_R5' else 'r2',person_score=.35,joint_score=.10)
 scenario=(FightScenario(FightSettings.parse(cfg),event_id_factory=lambda:f'mock-{next(ids)}',wall_ns=lambda:round(clock[0]*1e9)) if kind=='FIGHT_R5' else FallScenario(FallSettings.parse(cfg),event_id_factory=lambda:f'mock-{next(ids)}',wall_ns=lambda:round(clock[0]*1e9)))
 e=scenario.engine;e._fallback_admission=EventCooldown(10,clock=lambda:clock[0])
 if kind=='FIGHT_R5':e.display_policy.clock=lambda:clock[0]
 save(out/'settings.json',scenario.config.settings_payload())
 name='Fight_Review.mp4' if kind=='FIGHT_R5' else 'Fall_Review.mp4'
 video=out/name;assert not video.exists()
 container=av.open(str(video),'w');stream=container.add_stream('libx264',rate=10)
 stream.width=3200;stream.height=1800;stream.pix_fmt='yuv420p';stream.time_base=Fraction(1,90000)
 stream.codec_context.time_base=Fraction(1,90000);stream.codec_context.thread_count=2
 stream.options={'preset':'veryfast','crf':'21','bf':'0'}
 stats=dict(clip=key,engine=kind,observations=len(rows),current=0,held_only=0,events=0)
 diag=(out/'diagnostics.jsonl').open('x');pts=[]
 for row,m in zip(rows,maps):
  t=row['seconds'];clock[0]=t;w,h=row['width'],row['height']
  adjusted=[];people=[]
  for original,p in enumerate(row['persons']):
   box=[max(0.,min(float(limit),v)) for v,limit in zip(p['bbox'],[w,h,w,h])]
   if box!=p['bbox']:adjusted.append(dict(cache_order=original,before=p['bbox'],after=box))
   if box[2]<=box[0] or box[3]<=box[1]:continue
   people.append(dict(p,bbox=box,cache_order=original))
  row=dict(row,persons=people)
  fr=FrameEnvelope('offline',str(key),'review',key,1,row['sample'],round(t*1e9),1,w,h)
  ps=tuple(PoseObservation(fr,POSE_MODEL,p['score'],tuple(v/d for v,d in zip(p['bbox'],[w,h,w,h])),tuple((x/w,y/h,s) for x,y,s in p['keypoints'])) for p in row['persons'])
  tids=tuple(p['track_id'] for p in row['persons'])
  with patch('time.monotonic',lambda:t):displays,events=scenario.evaluate(fr,ps,tids)
  result=e.last_result;colors={};valid=0
  if kind=='FIGHT_R5':
   selected=[i for i,p in enumerate(ps) if inside(((p.bbox_xyxy[0]+p.bbox_xyxy[2])/2,p.bbox_xyxy[3]),scenario.config.roi)]
   valid=len(bodies(row,e.core.geometry)[0])
   for pair in result['pairs']:
    for i in pair['orders']:colors[selected[i]]='BLUE_CANDIDATE'
   for i in result['display_orders']:colors[selected[i]]='YELLOW_HOLD'
   for i in result['current_orders']:colors[selected[i]]='RED_CURRENT'
  else:
   colors={x['query']:x['state'] for x in result['colors']};valid=int(result['observable'])
  assert not events or result['hit']
  hold=any(x=='YELLOW_HOLD' for x in colors.values())
  stats['current']+=int(result['hit']);stats['held_only']+=int(hold and not result['hit']);stats['events']+=len(events)
  record=dict(sample=row['sample'],clip_frame=row['source_frame'],recording_frame=int(m['recording_frame']),pts=t,current=result['hit'],colors=colors,valid=valid,events=[x.event_id for x in events],notification=e.notification_status(),raw=result,bbox_adapter=adjusted,cache_orders=[p['cache_order'] for p in people])
  diag.write(json.dumps(record,default=plain)+'\n')
  img=Image.open(C/'clips'/key/'inputs/frames'/f"{row['sample']:04d}.png").convert('RGB').resize((2240,1260))
  overlay=Image.new('RGBA',img.size);d=ImageDraw.Draw(overlay)
  for i,p in enumerate(row['persons']):
   if p['score']<.35:continue
   b=[p['bbox'][0]*2240/w,p['bbox'][1]*1260/h,p['bbox'][2]*2240/w,p['bbox'][3]*1260/h]
   if b[2]<=b[0] or b[3]<=b[1]:continue
   state=colors.get(i);color=COLORS.get(state,(195,195,195))
   d.rectangle(b,fill=(*color,60) if state else None,outline=(*color,255),width=6 if state else 2)
   for a,z in ((5,6),(5,11),(6,12),(11,12),(5,7),(7,9),(6,8),(8,10),(11,13),(13,15),(12,14),(14,16)):
    pa,pb=p['keypoints'][a],p['keypoints'][z]
    if min(pa[2],pb[2])>=.1:d.line((pa[0]*2240/w,pa[1]*1260/h,pb[0]*2240/w,pb[1]*1260/h),fill=(225,225,225,255),width=3)
  img=Image.alpha_composite(img.convert('RGBA'),overlay).convert('RGB')
  canvas=Image.new('RGB',(3200,1800),(22,24,27));canvas.paste(img,(30,220));d=ImageDraw.Draw(canvas)
  d.text((30,35),kind+' | '+key,font=H,fill='white')
  lines=[f"Clip Frame {row['source_frame']}",f"Recording Frame {m['recording_frame']}",f'PTS {t:.3f} s','Person >= 0.35','Joint >= 0.10','Cooldown 10 s','Hold '+('3 s' if kind=='FIGHT_R5' else '1 s'),'CURRENT HIT' if result['hit'] else ('NO CURRENT HIT' if valid else 'INSUFFICIENT EVIDENCE'),'DISPLAY HOLD' if hold else 'No held display',f'New mock Event: {len(events)}',f"Event total: {stats['events']}"]
  for j,line in enumerate(lines):d.text((2340,240+j*75),line,font=F,fill='white')
  d.text((30,1560),'BLUE candidate | RED current | YELLOW display hold',font=H,fill='white')
  vf=av.VideoFrame.from_ndarray(np.asarray(canvas),format='rgb24');vf.pts=round(t*90000);vf.time_base=Fraction(1,90000);pts.append(t)
  for pkt in stream.encode(vf):container.mux(pkt)
  if row['sample']==0:canvas.save(out/'first_frame.png')
 for pkt in stream.encode():container.mux(pkt)
 container.close();diag.close()
 p=subprocess.run(['ffprobe','-v','error','-select_streams','v:0','-show_frames','-show_entries','frame=best_effort_timestamp_time','-of','json',str(video)],capture_output=True,timeout=600)
 assert p.returncode==0 and not p.stderr
 actual=[float(x['best_effort_timestamp_time']) for x in json.loads(p.stdout)['frames']]
 assert len(actual)==len(pts) and max(abs(a-b) for a,b in zip(actual,pts))<.00002
 stats.update(video=name,sha256=sha(video),pts_error=max(abs(a-b) for a,b in zip(actual,pts)),decode_errors=0)
 save(verify,stats);print(json.dumps(stats),flush=True);return stats
if __name__=='__main__':
 specs=json.loads((C/'specs.json').read_text());results=[]
 keys=list(specs)
 if '--first' in sys.argv:keys=keys[:1]
 for key in keys:
  for kind in ('FIGHT_R5','FALL_R2'):results.append(one(key,kind))
 save(R/'summary.json',results)