"""Standalone loopback benchmark driver. Python standard library only."""
import argparse,csv,json,sys,time,re,hashlib,ctypes,subprocess,struct
from pathlib import Path
from datetime import datetime,timezone

def dist(v):
    a=sorted(v)
    if not a:return {'n':0}
    return dict(n=len(a),p50_us=a[(len(a)-1)//2],p99_us=a[min(len(a)-1,int(len(a)*.99))],p999_us=a[min(len(a)-1,int(len(a)*.999))],max_us=a[-1],ge1000=sum(x>=1000 for x in a),ge5000=sum(x>=5000 for x in a),ge8000=sum(x>=8000 for x in a))

def analyze(folder):
    results={}
    for sender,receiver in [('gw','mt') ,('mt','gw')]:
        txt=(folder/f'{sender}.out').read_text()
        freq=int(re.search(r'qpc_freq=(\d+)',txt).group(1))
        send={}
        with (folder/f'{sender}-send.csv').open() as f:
            for r in csv.DictReader(f):
                k=(int(r['thread']),int(r['seq']));assert k not in send
                send[k]=(int(r['begin_qpc']),int(r['end_qpc']))
        lat=[];duration=[];postreturn=[];late=[];seen=set();completed_before_sender_return=0
        with (folder/f'{receiver}-rx.csv').open() as f:
            for r in csv.DictReader(f):
                k=(int(r['thread']),int(r['seq']));assert k not in seen;seen.add(k)
                start,done=send[k];assert start==int(r['begin_qpc'])
                end=int(r['end_qpc']);l=(end-start)*1e6/freq;s=(done-start)*1e6/freq;p=(end-done)*1e6/freq
                lat.append(l);duration.append(s);postreturn.append(p)
                if s>=5000 and end<done:completed_before_sender_return+=1
                if l>=1000:late.append(dict(thread=k[0],seq=k[1],start_qpc=start,send_end_qpc=done,rx_qpc=end,lat_us=l,send_us=s,after_send_us=p))
        assert seen==set(send)
        results[f'{sender}->{receiver}']={'latency':dist(lat),'send_duration_per_message':dist(duration),'after_send_return':dist(postreturn),'fast_send_under_500us_rx_ge5ms':sum(r['send_us']<500 and r['lat_us']>=5000 for r in late),'slow_send_ge5ms_but_peer_already_received':completed_before_sender_return,'late_samples':late}
    (folder/'analysis.json').write_text(json.dumps(results,indent=2))
    return {k:{a:b for a,b in v.items() if a!='late_samples'} for k,v in results.items()}

def utc():return datetime.now(timezone.utc).isoformat()
def counter():
    value=ctypes.c_longlong();freq=ctypes.c_longlong()
    assert ctypes.windll.kernel32.QueryPerformanceCounter(ctypes.byref(value)) and ctypes.windll.kernel32.QueryPerformanceFrequency(ctypes.byref(freq))
    return value.value,freq.value
def inspect(folder,state):
    endpoints={};pacing={}
    pattern=r'SOCKET index=(\d+) local=127.0.0.1:(\d+) remote=127.0.0.1:(\d+) tx=(\d+) rx=(\d+) sndbuf=(\d+) rcvbuf=(\d+) nodelay=(\d+)'
    for side in ['gw','mt']:
        stdout=(folder/(side+'.out')).read_text();rows=[tuple(map(int,m)) for m in re.findall(pattern,stdout)]
        assert len(rows)==(2 if state['split'] else 1)
        assert all(r[5:]==(65536,65536,1) for r in rows),rows
        assert sum(r[3] for r in rows)==sum(r[4] for r in rows)==1
        assert all(r[3]+r[4]==1 for r in rows) if state['split'] else rows[0][3:5]==(1,1)
        assert int(re.search(r'START qpc=(\d+)',stdout).group(1))==state['start_qpc']
        endpoints[side]=rows
        with (folder/(side+'-send.csv')).open() as f:samples=[{k:int(v) for k,v in r.items()} for r in csv.DictReader(f)]
        digest=hashlib.sha256()
        for r in sorted(samples,key=lambda x:(x['thread'],x['seq'])):digest.update(struct.pack('<IQQ',r['thread'],r['seq'],r['due_qpc']-state['start_qpc']))
        lateness=[(r['begin_qpc']-r['due_qpc'])*1e6/state['qpc_frequency'] for r in samples]
        assert all(v>=0 for v in lateness)
        pacing[side]={'n':len(samples),'schedule_sha256':digest.hexdigest(),'due_to_send_us':dist(lateness),'first_send_after_start_ms':(min(r['begin_qpc'] for r in samples)-state['start_qpc'])*1000/state['qpc_frequency'],'last_send_after_start_ms':(max(r['begin_qpc'] for r in samples)-state['start_qpc'])*1000/state['qpc_frequency']}
    for a,b in zip(endpoints['gw'],endpoints['mt']):assert a[1]==b[2] and a[2]==b[1] and a[3]==b[4] and a[4]==b[3]
    state['endpoints']=endpoints;state['pacing']=pacing
    return state
def run(name,iocp,secs):
    folder=OUT/name;folder.mkdir(parents=True,exist_ok=False);files=[];children=[]
    state={'name':name,'split':SPLIT,'iocp':iocp,'completion_workers_per_process':int(iocp),'seconds':secs,'status':'starting','started_utc':utc(),'exe_sha256':hashlib.sha256(EXE.read_bytes()).hexdigest(),'source_sha256':hashlib.sha256(EXE.with_suffix('.cpp').read_bytes()).hexdigest(),'rates':[1600,3200],'bursts':[4,8],'sender_threads_per_process':2,'receiver_threads_per_process':1,'seed':314159,'buffers':65536,'nodelay':1}
    try:
        for side,role,addr,rate,burst in [('gw','listen',str(PORT),1600,4),('mt','connect','127.0.0.1:'+str(PORT),3200,8)]:
            command=[str(EXE),role,addr,str(secs),str(rate),str(burst),'2','65536','1',side,'65536','0','-1','314159','0','0',str(int(SPLIT))]+([str(int(iocp))] if EXPERIMENT=='iocp' else [])
            out=(folder/(side+'.out')).open('w');err=(folder/(side+'.err')).open('w');files.extend([out,err])
            p=subprocess.Popen(command,cwd=folder,stdout=out,stderr=err,creationflags=subprocess.CREATE_NO_WINDOW);children.append(p);state[side+'_pid']=p.pid
            if side=='gw':time.sleep(.5)
        setup_deadline=time.monotonic()+15
        while not all((folder/(side+'-ready.txt')).exists() for side in ['gw','mt']):
            if any(p.poll() is not None for p in children):raise RuntimeError('Endpoint exited during setup')
            if time.monotonic()>setup_deadline:raise RuntimeError('Setup deadline exceeded')
            time.sleep(.025)
        now,freq=counter();state['qpc_frequency']=freq;state['start_qpc']=now+freq//2
        pending=folder/'start-qpc.tmp';pending.write_text(str(state['start_qpc']));pending.replace(folder/'start-qpc.txt')
        state['status']='running';(folder/'state.json').write_text(json.dumps(state,indent=2));print('RUNNING',name,flush=True)
        deadline=time.monotonic()+secs+25
        for p in children:p.wait(timeout=max(1,deadline-time.monotonic()))
        state['exit_codes']=[p.returncode for p in children]
        for f in files:f.close()
        assert state['exit_codes']==[0,0],state['exit_codes']
        assert all((folder/(side+'.err')).stat().st_size==0 for side in ['gw','mt'])
        state['results']=analyze(folder);inspect(folder,state)
        if EXPERIMENT=='iocp':inspect_completions(folder,state)
        state['status']='complete'
        print('COMPLETE',name,json.dumps({d:r['latency'] for d,r in state['results'].items()}),flush=True)
    except BaseException as e:state['status']='failed';state['error']=repr(e);raise
    finally:
        for p in children:
            if p.poll() is None:p.kill();p.wait()
        for f in files:
            if not f.closed:f.close()
        state['finished_utc']=utc();(folder/'state.json').write_text(json.dumps(state,indent=2))
    return state

def inspect_completions(folder,state):
    for side,peer in [('gw','mt'),('mt','gw')]:
        stdout=(folder/(side+'.out')).read_text()
        assert f"MODE iocp={int(state['iocp'])} producers=2 receiver_or_completion_workers=1" in stdout
        path=folder/(side+'-completion.csv')
        if not state['iocp']:
            assert not path.exists();continue
        with (folder/(side+'-send.csv')).open() as f:sends={(int(r['thread']),int(r['seq'])):(int(r['begin_qpc']),int(r['end_qpc'])) for r in csv.DictReader(f)}
        with (folder/(peer+'-rx.csv')).open() as f:rx={(int(r['thread']),int(r['seq'])):int(r['end_qpc']) for r in csv.DictReader(f)}
        durations=[];after_return=[];before_received=0;seen=set();pending=0
        with path.open() as f:
            for r in csv.DictReader(f):
                k=(int(r['thread']),int(r['seq']));assert k not in seen;seen.add(k)
                begin,returned=sends[k];completed=int(r['completion_qpc'])
                assert int(r['begin_qpc'])==begin and completed>=begin and int(r['bytes'])==32
                assert int(r['initiation_status']) in [0,997];pending+=int(r['initiation_status'])==997
                scale=1e6/state['qpc_frequency'];durations.append((completed-begin)*scale);after_return.append((completed-returned)*scale)
                before_received+=completed<rx[k]
        assert seen==set(sends)==set(rx)
        m=re.search(r'IOCP_AUDIT send_completions=(\d+) receive_completions=(\d+) high_water_pending=(\d+) immediate=(\d+) pending=(\d+)',stdout)
        assert m and int(m[1])==len(sends) and int(m[4])+int(m[5])==len(sends) and int(m[5])==pending
        state.setdefault('completions',{})[side]={'all_matched':True,'send_completion_observed':dist(durations),'completion_minus_api_return':dist(after_return),'completion_before_peer_observed':before_received,'high_water_pending':int(m[3]),'immediate_initiation':int(m[4]),'pending_initiation':pending,'receive_completions':int(m[2])}

def equal_schedules(states):
    for side in ['gw','mt']:
        assert len({s['pacing'][side]['schedule_sha256'] for s in states})==1,'Workload timing differs'
        assert len({s['pacing'][side]['n'] for s in states})==1,'Workload counts differ'

if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('mode',choices=['blocking','iocp','single','split'])
    parser.add_argument('--seconds',type=float,default=90)
    args=parser.parse_args()
    if not 0 < args.seconds <= 300:parser.error('--seconds must be >0 and <=300')
    EXPERIMENT='iocp' if args.mode in ['blocking','iocp'] else 'split'
    SPLIT=args.mode=='split'
    PORT=17798 if EXPERIMENT=='iocp' else 17796
    base=Path(__file__).resolve().parent
    EXE=base/('lbtest-'+EXPERIMENT+'.exe')
    if not EXE.exists():parser.error('Build both C++ programs first; see README.txt')
    OUT=base/'runs'
    name=datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%S%fZ')+'-'+args.mode
    state=run(name,args.mode=='iocp',args.seconds)
    print('SAVED',OUT/name,flush=True)