"""Build/run a shadow-only replay experiment, R288 before position 52.

No Candidate result is used by the live evaluator. Baseline metadata is captured
online; Reference executes on a deep copy of the actual current probe entry.
"""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import argparse,hashlib,json,os,shutil,time
import diagnose_context_access as access

ROOT=access.ROOT
HERE=Path(__file__).parent
BASE=ROOT/'test/context-access-20260917/worker'
OUT=ROOT/'test/context-shadow-20260917'
replace=access.replace


def generated_model(source,plan):
    a,b,e=access.dx.ex.function_span(source,'dx_model_eval_local_internal')
    fn=source[a:e]
    tail=fn[fn.index('double node_energy['):]
    generated=[]
    for pos in (16,52):
        operation=plan['code'][pos][0]
        expected=(45,0,3,4,64) if pos==16 else (152,28,44,40,178)
        qi,pipe,pi,gi,hi=expected
        expr=operation.split('=',1)[1].rstrip(';')
        for old,new in [(f'p[{pi}]','x[3]'),(f'g[{gi}].p','x[1]'),(f'g[{gi}].T','x[0]'),(f'h[{hi}]','x[2]')]:expr=expr.replace(old,new)
        generated.append(f'static double sr_original_{pos}(NativePropertyCache *properties,NativePipeCache *pipe_cache,const double *x){{return {expr};}}')
        header=f'''static int sr_tail_{pos}(double t,const double *y,SRContext *ctx,SRFrame *f){{
NativePropertyCache *properties=&ctx->context;NativePipeCache *pipe_cache=ctx->pipes;
double *p=f->p,*h=f->h,*q=f->q,*fb=f->fb,*w=f->w,*dy=f->dy;NativeGas *g=f->g;
(void)t;(void)y;(void)p;(void)h;(void)q;(void)fb;(void)w;(void)dy;(void)g;(void)properties;(void)pipe_cache;
'''
        generated.append(header+'\n'.join('\n'.join(c) for c in plan['code'][pos+1:])+'\n'+tail)
        hook=f'sr_shadow({pos},t,y,properties,pipe_cache,p,h,q,fb,g,w,sr_original_{pos},sr_tail_{pos});\n'
        # Both baseline capture and the live probe switch still execute the
        # original assignment. The hook returns void and cannot supply q.
        assert fn.count(operation)==3  # detailed baseline, timed baseline, probe
        fn=fn.replace(operation,hook+operation)
    return source[:a]+'\n'.join(generated)+'\n'+fn+source[e:]


def prepare():
    OUT.mkdir(exist_ok=True);work=OUT/'worker';work.mkdir(exist_ok=True)
    sources={p.name:p.read_text(encoding='utf-8') for p in BASE.glob('*.c')}
    assert sources
    before={k:hashlib.sha256(v.encode()).hexdigest() for k,v in sources.items()}
    certified=json.loads((BASE/'build.json').read_text(encoding='utf-8'))['instrumentedHashes']
    assert before==certified, 'Access-validated source changed; revalidate it before extending the allowlist.'
    plan=json.loads((ROOT/'test/context-fallback-20260917/plan.json').read_text(encoding='utf-8'))
    sources['model.c']=generated_model(sources['model.c'],plan)
    sources['model.c']=replace(sources['model.c'],'ax_result(result);dx_eval_end();','sr_eval_result(result,dy,w);dx_eval_end();')
    # Every native entry reachable by either target is guarded during replay.
    # Candidate's semantic interpreter has no calls to these physics functions.
    native_functions={
        'properties.c':['property_pt','property_density','property_viscosity','native_temperature_ph_context','local_isentropic','isentropic','state_valve','native_density','native_temperature_ph','native_viscosity'],
        'pipe.c':['native_pipe_flow_cached_context','native_pipe_flow_context','native_pipe_resistance'],
    }
    for module,names in native_functions.items():
        for name in names:sources[module]=access.change_function(sources[module],name,lambda body:'sr_native_enter(__func__);\n'+body)
    del sources['context_access_diag.c']
    sources['context_shadow_replay.c']=(HERE/'context_shadow_replay.c').read_text(encoding='utf-8')
    for p in BASE.glob('*.h'):shutil.copyfile(p,work/p.name)
    shutil.copyfile(HERE/'context_shadow_replay.h',work/'context_shadow_replay.h')
    header=(work/'context_access_diag.h').read_text(encoding='utf-8')
    header=header.replace('#include "kernels.h"','#include "kernels.h"\n#include "context_shadow_replay.h"')
    header=replace(header,'(x)|=(v); ax_access("write",&(x),sizeof(x),#x,__func__);','(x)|=(v); sr_or(&(x),(v),#x,__func__);')
    (work/'context_access_diag.h').write_text(header,encoding='utf-8')
    cc,sun,_=access.dx.ex.builder.toolchain();flags,libs,dlls,exe=access.dx.ex.builder.platform_build_inputs(sun)
    flags+=['-DLP_OBSERVE=0'];started=time.perf_counter()
    def compile_one(item):
        name,code=item;path=work/name;path.write_text(code,encoding='utf-8',newline='\n');obj=path.with_suffix('.o');log=[]
        access.dx.ex.builder._command([cc,*flags,'-I',str(work),'-I',str(access.dx.ex.builder.NATIVE/'include'),'-I',str(sun/'include'),'-c',str(path),'-o',str(obj)],log=log,timeout=240)
        return obj,log
    with ThreadPoolExecutor(max_workers=4) as pool:objects=list(pool.map(compile_one,sources.items()))
    log=[];access.dx.ex.builder._command([cc,*flags,*[str(o) for o,_ in objects],*access.dx.ex.builder.link_library_arguments(libs),'-lm','-o',str(work/exe)],log=log)
    for dll in dlls:shutil.copyfile(dll,work/dll.name)
    (work/'build.log').write_text('\n'.join(sum([v for _,v in objects],[])+log),encoding='utf-8')
    guard_before=access.dx.function(sources['local_probe_support.c'],'lp_reuse')
    guard_certified=access.dx.function((BASE/'local_probe_support.c').read_text(encoding='utf-8'),'lp_reuse')
    assert guard_before==guard_certified
    access.dx.ex.write(work/'build.json',dict(sourceHashes=before,instrumentedHashes={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in work.glob('*.c')},seconds=time.perf_counter()-started,guardedNativeFunctions=native_functions,wholeContextGuardUnchanged=True))
    print('BUILT independent shadow worker',flush=True)


def run(stage):
    worker_hash=hashlib.sha256((OUT/'worker/model.exe').read_bytes()).hexdigest()
    if stage==52:
        gate=json.loads((OUT/'r288/shadow-summary.json').read_text(encoding='utf-8'))
        assert gate['total']==896 and gate['accepted']==896 and gate['mismatches']==0 and gate['liveContamination']==0, 'R288 must pass first.'
        assert gate['validatedWorkerSha256']==worker_hash, 'R288 must validate this exact worker build first.'
    os.environ['CONTEXT_SHADOW_POSITION']=str(stage)
    access.dx.OUT=OUT
    label='r288' if stage==16 else 'position52'
    access.dx.run(label,mode=1,matrices=True)
    summary=json.loads((OUT/label/'shadow-summary.json').read_text(encoding='utf-8'))
    assert summary['total']==896 and summary['mismatches']==0 and summary['liveContamination']==0
    summary['validatedWorkerSha256']=worker_hash
    access.dx.ex.write(OUT/label/'shadow-summary.json',summary)
    print(json.dumps(summary,indent=2),flush=True)


if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('action',choices=['prepare','run']);p.add_argument('--position',type=int,choices=[16,52],default=16);a=p.parse_args()
    if a.action=='prepare':prepare()
    else:run(a.position)