"""G: fixed video frames, independently paced sender/receiver, real WebSocket server.
Run against the already-running local server; no source frame/video is copied.
"""
from pathlib import Path
import sys
import os
import json
import time
import asyncio
import statistics
import subprocess
import hashlib
import argparse
import resource
from collections import defaultdict

ROOT=Path(__file__).resolve().parents[3]
sys.path.insert(0,str(ROOT/'code'))
import cv2
from websockets.asyncio.client import connect
from stream_client import FrameSource,encode
OUT=Path(__file__).resolve().parent/'evidence'


def q(values,p):
    if not values:return None
    return round(sorted(values)[min(len(values)-1,int((len(values)-1)*p))],2)


def input_frames(source, fps, duration=5):
    if source=='image6':
        view=FrameSource(str(ROOT/'add/image6.jpg'),width=1280,pan_step=5.6)
        out=[]
        for n in range(int(duration*fps)):
            out.append((encode(view.frame(n),1280,75),view.offset,view.scale))
        view.release()
        return out
    cap=cv2.VideoCapture(str(ROOT/source))
    sample=[]; frame_count=int(cap.get(cv2.CAP_PROP_FRAME_COUNT)); native_fps=cap.get(cv2.CAP_PROP_FPS)
    start=max(0,(frame_count-int(duration*native_fps))//2)
    for n in range(int(duration*fps)):
        cap.set(cv2.CAP_PROP_POS_FRAMES,start+int(n*native_fps/fps))
        ok,frame=cap.read()
        if not ok:break
        sample.append((encode(frame,1280,75),(0,0),1.))
    cap.release()
    return sample


async def trial(source,frames,fps,rate,trial_id):
    # Real-time source sampling is frozen across the two pace arms. At .5x each
    # source frame is shown longer, rather than processing a different sequence.
    secret=(ROOT/'.demo-token').read_text().strip()
    started=time.perf_counter(); results=[]; sent={}; first={}; transitions=defaultdict(list); errors=[]
    max_tracks=0; pid=int(subprocess.check_output(['lsof','-tiTCP:8099','-sTCP:LISTEN'],text=True).splitlines()[0])
    from psutil import Process
    proc=Process(pid); cpu0=proc.cpu_times(); cpu_before=time.perf_counter()
    async with connect('ws://127.0.0.1:8099/v1/stream',additional_headers={'Authorization':f'Bearer {secret}'},max_size=16*1024*1024) as socket:
        hello=json.loads(await socket.recv())
        if hello.get('type')!='hello':raise RuntimeError('server did not greet authenticated session')
        await socket.send(json.dumps({'type':'config','tags':False,'empty_slots':False}))
        await socket.recv()
        async def receive():
            nonlocal max_tracks
            while True:
                try: raw=await asyncio.wait_for(socket.recv(), timeout=2)
                except asyncio.TimeoutError: break
                msg=json.loads(raw)
                if msg.get('type')!='result':
                    errors.append(msg);continue
                now=time.perf_counter(); tracks=msg['tracks']; max_tracks=max(max_tracks,len(tracks))
                visible=[]
                for tr in tracks:
                    state=tr['identity']; sku=state.get('sku_id'); tid=tr['track_id']
                    if sku:
                        if tid not in first:first[tid]=now
                        previous=transitions[tid]
                        if not previous or previous[-1]['sku']!=sku:previous.append({'ms':round((now-started)*1000,1),'sku':sku,'state':state.get('state')})
                    box=tr.get('box') or tr.get('raw_box')
                    origin,scale=frames[min(len(frames)-1,max(0,msg['frame_id']-1))][1:]
                    visible.append({'id':tid,'sku':sku,'state':state.get('state'),'identity_level':state.get('identity_level'),'box':[round(box[j]/scale+origin[j%2],1) for j in range(4)] if box else [],'object_generation':tr.get('object_generation'),'restarts':tr.get('restarts'), 'source':state.get('source')})
                results.append({'frame_id':msg['frame_id'],'received_ms':round((now-started)*1000,1),'sent_ms':round((sent.get(msg['frame_id'],started)-started)*1000,1),'timing_ms':msg.get('timing_ms'),'stats':msg.get('stats'), 'tracks':visible})
        receiver=asyncio.create_task(receive())
        for n,(frame,*_) in enumerate(frames):
            deadline=started+n/(fps*rate)
            await asyncio.sleep(max(0,deadline-time.perf_counter()))
            await socket.send((n+1).to_bytes(4,'big')+frame)
            sent[n+1]=time.perf_counter()
        await asyncio.wait_for(receiver, timeout=max(15,len(frames)/fps/rate+6))
    elapsed=time.perf_counter()-started; cpu1=proc.cpu_times();cpu=(cpu1.user+cpu1.system-cpu0.user-cpu0.system)
    received={r['frame_id'] for r in results}
    stages={k:[r['timing_ms'][k] for r in results if r.get('timing_ms') and r['timing_ms'].get(k) is not None] for k in ['decode','detect','track','total']}
    rtt=[r['received_ms']-r['sent_ms'] for r in results]
    levels=defaultdict(int)
    for r in results:
        for t in r['tracks']:
            if t['sku']: levels[t.get('identity_level') or 'unset']+=1
    summary={'source':source,'fps':fps,'rate':rate,'trial':trial_id,'input_frames':len(frames),'sent':len(sent),'processed':len(results),'not_returned':len(sent)-len(received),'elapsed_s':round(elapsed,2),'server_cpu_s':round(cpu,2),'server_cpu_pct_one_core':round(cpu/max(.001,time.perf_counter()-cpu_before)*100,1),'rtt_ms':{'p50':q(rtt,.5),'p95':q(rtt,.95)},'stages_ms':{k:{'p50':q(v,.5),'p95':q(v,.95)} for k,v in stages.items()},'max_tracks':max_tracks,'ever_labelled_track_ids':len(first),'first_visible_ms':min((v-started)*1000 for v in first.values()) if first else None,'label_flips':sum(max(0,len(x)-1) for x in transitions.values()),'identity_levels':dict(levels),'errors':errors,'latest_stats':results[-1]['stats'] if results else None}
    output={'schema_version':2,'coordinate_basis':'displayed box mapped with result frame origin','summary':summary,'timelines':results,'transitions':dict(transitions),'input_hash':hashlib.sha256(b''.join(x[0] for x in frames)).hexdigest()}
    path=OUT/f'replay-{Path(source).stem}-{fps}fps-{rate}x-{trial_id}.json';path.write_text(json.dumps(output,ensure_ascii=False,indent=1))
    print(path.name,summary['processed'],'/',summary['sent'],'first',summary['first_visible_ms'],'detect',summary['stages_ms']['detect'],'cpu',summary['server_cpu_s'],flush=True)


async def main():
    parser=argparse.ArgumentParser();parser.add_argument('--source',default='image6');parser.add_argument('--fps',type=int,default=8);parser.add_argument('--seconds',type=int,default=5);parser.add_argument('--trials',type=int,default=2)
    args=parser.parse_args();OUT.mkdir(exist_ok=True)
    frames=input_frames(args.source,args.fps,args.seconds)
    if not frames:raise SystemExit('source yielded no frames')
    for i in range(args.trials):
        for rate in (1,.5):await trial(args.source,frames,args.fps,rate,i)

if __name__=='__main__':asyncio.run(main())