"""Recorded-file Pose only; independent Triton and no production network access."""
import csv
import hashlib
import json
import pathlib
import subprocess
import sys
import time
from types import SimpleNamespace
import urllib.request

import av
import cv2
import numpy as np

sys.path.insert(0,'/opt/contextra')
from app.rfdetr_postprocess import parse_boxes_labels, parse_pose_keypoints
from app.byte_tracking import ByteTracker, TrackingSettings
from app.pose_nms import pose_nms_indices

ROOT=pathlib.Path('/eval')
assert not (ROOT/'inputs/pose.json').exists(), 'Existing Pose cache must be validated, not overwritten'
frames_dir=ROOT/'inputs/frames'
frames_dir.mkdir(parents=True,exist_ok=True)
config='''name: "RFDETRKeypointPreview"
platform: "tensorrt_plan"
max_batch_size: 8
version_policy { specific { versions: [8] } }
input [{ name: "input" data_type: TYPE_FP16 dims: [3,576,576] }]
output [
{ name: "dets" data_type: TYPE_FP32 dims: [100,4] },
{ name: "labels" data_type: TYPE_FP32 dims: [100,2] },
{ name: "keypoints" data_type: TYPE_FP32 dims: [100,34,8] }]
instance_group [{ kind: KIND_GPU count: 1 gpus: [0] }]
'''
(ROOT/'pose_models/RFDETRKeypointPreview/config.pbtxt').write_text(config)
rows=list(csv.DictReader(open('/eval/inputs/source_map.csv')))
plan=json.loads((ROOT/'inputs/sampling_plan.json').read_text())
pts=np.array([float(r['elapsed_seconds']) for r in rows])
grid=np.asarray(plan['target_seconds'],dtype=np.float64)
indices=np.array([int(np.argmin(abs(pts-t))) for t in grid])
assert len(set(indices))==len(grid) and np.all(np.diff(indices)>0)
assert (indices+1).tolist()==plan['source_frame_indices']
targets={int(k+1):j for j,k in enumerate(indices)}
tracker=ByteTracker(TrackingSettings(frame_rate=30))
report=dict(status='STARTED',fps=10,sampling='nearest actual PTS to 0.1s grid; unique source frames',
    person_threshold=.15,joint_threshold=.3,nms_iou=.85,
    tracking=tracker.settings.payload(),
    preprocess='CPU OpenCV bilinear stretch RGB576, ImageNet mean/std, FP32 arithmetic -> FP16 NCHW; GPU resize bit-equivalence not claimed',
    model='RFDETRKeypointPreview',engine_sha256=hashlib.sha256((ROOT/'pose_models/RFDETRKeypointPreview/8/model.plan').read_bytes()).hexdigest(),
    production_queue_used=False,experiment_instances=1,source_files={})
for name in ('rfdetr_postprocess.py','byte_tracking.py','tracking_kernel.py','pose_nms.py'):
    src=pathlib.Path('/opt/contextra/app')/name
    dest=ROOT/'reports'/('source_'+name)
    dest.write_bytes(src.read_bytes())
    report['source_files'][name]=hashlib.sha256(src.read_bytes()).hexdigest()
log=open(ROOT/'reports/isolated_triton.log','w')
proc=subprocess.Popen(['/opt/tritonserver/bin/tritonserver','--model-repository=/eval/pose_models',
    '--model-control-mode=explicit','--load-model=RFDETRKeypointPreview','--allow-grpc=false',
    '--allow-metrics=false','--http-thread-count=2','--pinned-memory-pool-byte-size=67108864',
    '--cuda-memory-pool-byte-size=0:67108864'],stdout=log,stderr=subprocess.STDOUT)
start=time.monotonic()
all_rows=[]
raw_cache=open(ROOT/'inputs/pose_candidates_prefilter.jsonl','x')
report['candidate_cache']='all decoded queries before person threshold, NMS and Tracking; original 17 joint scores retained'
try:
    for _ in range(90):
        if proc.poll() is not None: raise RuntimeError('Isolated Triton exited')
        try:
            with urllib.request.urlopen('http://127.0.0.1:8000/v2/models/RFDETRKeypointPreview/versions/8/ready',timeout=5) as r:
                if r.status==200: break
        except Exception: pass
        time.sleep(1)
    else: raise TimeoutError('Isolated Pose READY timeout')
    source=av.open('/source.mp4')
    source.streams.video[0].codec_context.thread_count=2
    for number,frame in enumerate(source.decode(video=0),1):
        if number>max(targets): break
        if number not in targets: continue
        sample=targets[number]
        bgr=frame.to_ndarray(format='bgr24')
        height,width=bgr.shape[:2]
        rgb=cv2.cvtColor(cv2.resize(bgr,(576,576),interpolation=cv2.INTER_LINEAR),cv2.COLOR_BGR2RGB)
        x=(rgb.astype(np.float32)/255-np.array([.485,.456,.406],np.float32))/np.array([.229,.224,.225],np.float32)
        x=np.ascontiguousarray(x.transpose(2,0,1)[None],dtype=np.float16)
        header=json.dumps(dict(inputs=[dict(name='input',shape=list(x.shape),datatype='FP16',
            parameters=dict(binary_data_size=x.nbytes))],outputs=[dict(name=n,parameters=dict(binary_data=True))
            for n in ('dets','labels','keypoints')])).encode()
        req=urllib.request.Request('http://127.0.0.1:8000/v2/models/RFDETRKeypointPreview/versions/8/infer',
            data=header+x.tobytes(),headers={'Content-Type':'application/octet-stream',
            'Inference-Header-Content-Length':str(len(header))})
        with urllib.request.urlopen(req,timeout=30) as r:
            n=int(r.headers['Inference-Header-Content-Length']); payload=r.read()
        meta=json.loads(payload[:n]); raw={}; offset=n
        for out in meta['outputs']:
            size=out['parameters']['binary_data_size']; assert out['datatype']=='FP32'
            raw[out['name']]=np.frombuffer(payload[offset:offset+size],dtype='<f4').reshape(out['shape'])
            offset+=size
        assert offset==len(payload)
        dets=raw['dets'][0]; labels=raw['labels'][0]; kp=raw['keypoints'][0]
        parsed=parse_boxes_labels(dets,labels,threshold=.1)
        kps=parse_pose_keypoints(kp,parsed['model_classes'])
        observations=[]
        for q in np.flatnonzero(parsed['keep']):
            cx,cy,bw,bh=map(float,dets[q]); box=(cx-bw/2,cy-bh/2,cx+bw/2,cy+bh/2)
            observations.append(SimpleNamespace(class_id=0,score=float(parsed['scores'][q]),bbox_xyxy=box,model_key='RFDETRKeypointPreview',
                frame=SimpleNamespace(frame_width=width,frame_height=height),query=int(q)))
        selected=pose_nms_indices(observations,.85)
        observations=[observations[i] for i in selected]
        ids=tracker.update(observations)
        tracker_ids={o.query:tid for o,tid in zip(observations,ids)}
        nms_queries={o.query for o in observations}
        candidates=[]
        for q in range(len(dets)):
            cx,cy,bw,bh=map(float,dets[q])
            points=kps[q].copy();points[:,0]*=width;points[:,1]*=height
            candidates.append(dict(query=q,score=float(parsed['scores'][q]),
                bbox=[(cx-bw/2)*width,(cy-bh/2)*height,(cx+bw/2)*width,(cy+bh/2)*height],
                keypoints=points.tolist(),parse_keep=bool(parsed['keep'][q]),
                nms_keep=q in nms_queries,track_id=tracker_ids.get(q)))
        raw_cache.write(json.dumps(dict(sample=sample,source_frame=number,seconds=float(pts[number-1]),
            width=width,height=height,candidates=candidates),separators=(',',':'),allow_nan=False)+'\n')
        persons=[]
        for o,tid in zip(observations,ids):
            if o.score < .15: continue
            points=kps[o.query].copy(); points[:,0]*=width; points[:,1]*=height
            persons.append(dict(track_id=tid,query=o.query,score=o.score,
                bbox=(np.array(o.bbox_xyxy)*[width,height,width,height]).tolist(),keypoints=points.tolist(),
                valid_joints=int((points[:,2]>=.3).sum())))
        path=frames_dir/(f'{sample:04d}.png')
        assert cv2.imwrite(str(path),bgr,[cv2.IMWRITE_PNG_COMPRESSION,1])
        all_rows.append(dict(sample=sample,target_seconds=float(grid[sample]),source_frame=number,
            seconds=float(pts[number-1]),width=width,height=height,persons=persons))
        if sample%50==0:
            print(json.dumps(dict(sample=sample,people=len(persons),wall=round(time.monotonic()-start,1))),flush=True)
            mem=int(subprocess.check_output(['nvidia-smi','--query-gpu=memory.free','--format=csv,noheader,nounits']).decode().strip())
            if mem<4096: raise RuntimeError('GPU safety margin reached')
        time.sleep(.02)
    source.close()
    assert len(all_rows)==len(grid)
    assert [r['source_frame'] for r in all_rows]==plan['source_frame_indices']
    (ROOT/'inputs/pose.json').write_text(json.dumps(all_rows))
    with (ROOT/'inputs/sampling.csv').open('w') as f:
        w=csv.DictWriter(f,fieldnames=['sample','target_seconds','source_frame','seconds'])
        w.writeheader(); w.writerows({k:r[k] for k in w.fieldnames} for r in all_rows)
    report.update(status='PASS',samples=len(all_rows),people_min=min(len(r['persons']) for r in all_rows),
        people_max=max(len(r['persons']) for r in all_rows),
        untracked_observations=sum(p['track_id'] is None for r in all_rows for p in r['persons']),
        max_sampling_error_seconds=float(max(abs(pts[indices]-grid))))
finally:
    raw_cache.close()
    proc.terminate()
    try: proc.wait(timeout=20)
    except subprocess.TimeoutExpired: proc.kill(); proc.wait(timeout=10)
    log.close()
    report['wall_seconds']=round(time.monotonic()-start,3)
    (ROOT/'reports/pose_summary.json').write_text(json.dumps(report,indent=2))
    print(json.dumps(report),flush=True)