"""Only group 6 / region 288 / position 16, independent real-skip workers."""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import argparse,hashlib,json,os,re,shutil,subprocess,time
import diagnose_context_shadow as shadow

ROOT=shadow.ROOT;HERE=Path(__file__).parent
BASE=ROOT/'test/local-probe-20260917/worker'
ACCESS=ROOT/'test/context-access-20260917/worker'
OUT=ROOT/'test/r288-real-skip-20260917'
ex=shadow.access.dx.ex
replace=shadow.replace


def write(path,obj):ex.write(path,obj)


def fn(source,name):
    if name=='ax_test':
        start=source.index('unsigned ax_test(');end=source.index('\nvoid ax_scalar',start)
        return source[start:end]
    if name=='event':
        start=source.index('static Event *event(');end=source.index('\nstatic int state_pointer',start)
        return source[start:end]
    return shadow.access.dx.function(source,name)


def core(audit):
    source=(HERE/'context_shadow_replay.c').read_text(encoding='utf-8')
    certified=json.loads((shadow.OUT/'worker/build.json').read_text(encoding='utf-8'))['instrumentedHashes']['context_shadow_replay.c']
    assert hashlib.sha256(source.encode()).hexdigest()==certified, 'Use the shadow-certified source unchanged.'
    prefix=source[:source.index('typedef struct {\n    unsigned long long total')]
    prefix=prefix.replace('static void save_environment','static __attribute__((unused)) void save_environment').replace('static void restore_environment','static __attribute__((unused)) void restore_environment')
    globals='''
static Plan plan;
static int phase,position=16,group,configured=16;
static unsigned long long jac;
static double sim_time,op_inputs[4];
static NativePropertyCache *bound;
static NativePipeCache *bound_pipes;
static SRContext overlay;
static size_t probe_entry_count;
static void fatal(const char *s){fprintf(stderr,"real-skip fatal: %s\\n",s);abort();}
static void initialize(void){}
void sr_native_enter(const char *name){(void)name;if(phase==CANDIDATE)fatal("Candidate called physics");}
'''
    names=['locate','event','state_pointer','floating_field','ax_access','sr_or','ax_bind','ax_begin','ax_query','ax_match','ax_new','ax_test','ax_scalar','path_of_plan','ax_end','ax_result','from_live','same_medium_key','first_match','map_slot','scalar_lookup','state_read_equal','replay_overlay']
    parts=[prefix,globals]+[fn(source,n) for n in names]
    result='\n'.join(parts).replace('(size_t)slot<entry.context.count','(size_t)slot<probe_entry_count')
    assert fn(result,'replay_overlay')==fn(source,'replay_overlay').replace('(size_t)slot<entry.context.count','(size_t)slot<probe_entry_count')
    if audit:
        start=source.index('typedef struct {const char *name;size_t offset,size;} Field;')
        end=source.index('static void dump_blob(',start)
        comparison=source[start:end]
        a,b,e=ex.function_span(comparison,'compare_frame');comparison=comparison[:a]+comparison[e:]
        result+='\n'+comparison
    return result


def model_source(source,plan):
    a,b,e=ex.function_span(source,'model_eval_local_internal');body=source[a:e]
    op=plan['code'][16][0]
    assert body.count(op)==2 and plan['regions'][288]==[16,17]
    expr=op.split('=',1)[1].rstrip(';')
    for old,new in [('p[3]','x[3]'),('g[4].p','x[1]'),('g[4].T','x[0]'),('h[64]','x[2]')]:expr=expr.replace(old,new)
    helper=f'static double rr_original16(NativePropertyCache *properties,NativePipeCache *pipe_cache,const double *x){{return {expr};}}\n'
    ins='(double[]){g[4].T,g[4].p,h[64],p[3]}'
    baseline='rr_record_begin(properties,pipe_cache,'+ins+');\n#if RR_SKIP\n'+op.replace('native_pipe_flow_cached_context','rrrec_native_pipe_flow_cached_context').replace('native_temperature_ph_context','rrrec_native_temperature_ph_context')+'\n#else\n'+op+'\n#endif\nrr_record_end(q[45]);'
    body=body.replace(op,baseline,1)
    # The case is reachable only after the original whole-context guard failed.
    start=body.index('case 16:{');end=body.index('break;}',start)
    case=body[start:end]
    case=replace(case,op,'if(lp_color==6 && region==288){q[45]=rr_execute(properties,pipe_cache,'+ins+',rr_original16);}else{'+op+'}')
    body=body[:start]+case+body[end:]
    # Audit every Jacobian baseline/probe's exit, including all other groups.
    body=replace(body,'return 1;}','rr_eval_exit(t,y,dy,w,1,properties,pipe_cache,jacobian);return 1;}')
    return source[:a]+helper+body+source[e:]


def prepare(audit,skip):
    if not audit:
        gate=json.loads((OUT/'audit-skip-run/validation.json').read_text(encoding='utf-8'))
        assert gate['exact'] and gate['counters']['realSkips']==896
    OUT.mkdir(exist_ok=True)
    label=('audit' if audit else 'perf')+('-skip' if skip else '-control')
    work=OUT/label;work.mkdir(exist_ok=True)
    sources={p.name:p.read_text(encoding='utf-8') for p in BASE.glob('*.c')}
    original_hashes={k:hashlib.sha256(v.encode()).hexdigest() for k,v in sources.items()}
    plan=json.loads((ROOT/'test/context-fallback-20260917/plan.json').read_text(encoding='utf-8'))
    sources['model.c']=model_source(sources['model.c'],plan)
    sources['common.c']=replace(sources['common.c'],'lp_start();','lp_start();rr_start();')
    sources['common.c']=replace(sources['common.c'],'lp_finish();','lp_finish();rr_finish();')
    sources['cvode_solver.c']=replace(sources['cvode_solver.c'],'{lp_color=-1;uint64_t start=lp_tick();','{lp_color=-1;rr_jacobian();uint64_t start=lp_tick();')
    sources['cvode_solver.c']=replace(sources['cvode_solver.c'],'if(!result)lp_matrix(t,N_VGetArrayPointer(y),SUNDenseMatrix_Data(matrix));','if(!result){lp_matrix(t,N_VGetArrayPointer(y),SUNDenseMatrix_Data(matrix));rr_matrix(t,N_VGetArrayPointer(y),SUNDenseMatrix_Data(matrix));}')
    sources['pipe.c']=shadow.access.change_function(sources['pipe.c'],'native_pipe_flow_cached_context',lambda body:'if(rr_replaying)abort();if(rr_original_scope)rr_native_calls++;\n'+body)
    # Original physics files are NOT access-instrumented. A separate namespace
    # is called only to record the baseline target operation's required IR.
    recording={name:(ACCESS/name).read_text(encoding='utf-8') for name in ['properties.c','pipe.c']}
    names=set()
    for code in recording.values():
        names.update(re.findall(r'(?m)^(?:static\s+|NATIVE_COMPONENT_INTERNAL\s+)?(?:int|double|void|NativePropertyState\s*\*|NativeGas|NativeJacobianScalarEntry\s*\*)\s*(\w+)\([^;{}]*\)\s*\{',code))
    names.add('helium_medium')
    aliases='\n'.join('#define '+n+' rrrec_'+n for n in sorted(names))+'\n'
    for name,code in recording.items():sources['record_'+name]=aliases+code
    sources['real_skip_runtime.c']=(HERE/'real_skip_runtime.c').read_text(encoding='utf-8')
    for name in ['model.h','local_probe.h']:shutil.copyfile(BASE/name,work/name)
    shutil.copyfile(ACCESS/'context_fallback_diag.h',work/'context_fallback_diag.h')
    shutil.copyfile(HERE/'context_shadow_replay.h',work/'context_shadow_replay.h')
    shutil.copyfile(HERE/'real_skip.h',work/'real_skip.h')
    header=(ACCESS/'context_access_diag.h').read_text(encoding='utf-8').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')
    (work/'real_skip_core.inc').write_text(core(audit),encoding='utf-8',newline='\n')
    cc,sun,_=ex.builder.toolchain();flags,libs,dlls,exe=ex.builder.platform_build_inputs(sun)
    flags+=['-DLP_OBSERVE=0',f'-DRR_AUDIT={int(audit)}',f'-DRR_SKIP={int(skip)}'];started=time.perf_counter()
    def compile_one(item):
        name,code=item;path=work/name
        prefix='#ifdef _WIN32\n#ifndef _WIN32_WINNT\n#define _WIN32_WINNT 0x0600\n#endif\n#endif\n#include "real_skip.h"\n#include <stdlib.h>\n'
        if name.startswith('record_'):
            code=code[len(aliases):];prefix=aliases+prefix
        path.write_text(prefix+code,encoding='utf-8',newline='\n');obj=path.with_suffix('.o');log=[]
        ex.builder._command([cc,*flags,'-I',str(work),'-I',str(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=[];ex.builder._command([cc,*flags,*[str(p) for p,_ in objects],*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')
    assert sources['local_probe_support.c']==(BASE/'local_probe_support.c').read_text(encoding='utf-8')
    write(work/'build.json',dict(audit=audit,skip=skip,seconds=time.perf_counter()-started,originalHashes=original_hashes,guardUnchanged=True,guardMatchesShadow=True,sourceHashes={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in work.glob('*.c')},coreSha256=hashlib.sha256((work/'real_skip_core.inc').read_bytes()).hexdigest()))
    print('BUILT',label,flush=True)


def run(audit,skip,label=None,force=0):
    worker=('audit' if audit else 'perf')+('-skip' if skip else '-control')
    label=label or worker+'-run';folder=OUT/label;folder.mkdir(exist_ok=True)
    env=os.environ.copy();env['LOCAL_PROBE_MASK']='0x7ffffff'
    if audit and skip:env['RR_EXPECTED']=str(OUT/'audit-control-run/audit.bin')
    if force:env['RR_FORCE_REJECT_EVERY']=str(force)
    args=[str(OUT/worker/'model.exe'),'--method','BDF','--start','0','--stop','10','--sample-step','.01','--max-step','1e30','--rtol','1e-8','--timeout','300','--sample-file',str(folder/'states.bin'),'--output-block-file',str(folder/'outputs.bin'),'--output',str(folder/'result.json')]
    start=time.perf_counter()
    with (folder/'stderr.log').open('wb') as f:p=subprocess.run(args,cwd=folder,env=env,stdout=subprocess.PIPE,stderr=f,timeout=330,creationflags=subprocess.CREATE_NO_WINDOW)
    assert p.returncode==0,(label,p.returncode,(folder/'stderr.log').read_text()[-4000:])
    r=json.loads((folder/'result.json').read_text());d=json.loads((folder/'probe.json').read_text());c=json.loads((folder/'real-skip.json').read_text())
    ref=json.loads((BASE.parent/'all-run-0/measurement.json').read_text(encoding='utf-8'))
    keys=['success','finalState','final','propertyWarnings','acceptedSteps','rejectedSteps','stateTransitions','solverStarts','nfev','njev','nlu']
    assert all(r[k]==ref[k] for k in keys),[(k,r[k],ref[k]) for k in keys if r[k]!=ref[k]]
    for key in ['newtonIterations','newtonConvergenceFailures','modelCalls','groups','contextCopiedBytes','contextComparedBytes']:assert d[key]==ref['diagnostic'][key],key
    hashes={}
    for name in ['states','outputs','events']+(['jacobians'] if audit else []):
        with (folder/(name+'.bin')).open('rb') as f:hashes[name]=hashlib.file_digest(f,'sha256').hexdigest()
        if name=='jacobians':
            with (BASE.parent/'all-audit/jacobians.bin').open('rb') as f:assert hashes[name]==hashlib.file_digest(f,'sha256').hexdigest()
        else:assert hashes[name]==ref[name+'Sha256'],name
    expected_rejects=(896+force-1)//force if force else 0
    assert c['targetVisits']==896 and c['referenceExecutions']==0 and c['mismatches']==0
    assert c['candidateAttempts']==(896 if skip else 0)
    assert c['realSkips']==(896-expected_rejects if skip else 0)
    assert c['nativeOriginalExecutions']==(expected_rejects if skip else 896)
    assert c['commits']==c['semanticSuccess']==c['realSkips']
    assert c['fallbackExecutions']==c['rejects']==expected_rejects
    if audit:assert c['evaluationAudits']==896*28 and c['targetEntryAudits']==c['targetExitAudits']==896
    record=dict(label=label,audit=audit,skip=skip,force=force,exact=True,solveSeconds=r['solveSeconds'],solveCpuSeconds=r['solveCpuSeconds'],jacobianSeconds=d['jacobianSeconds'],processSeconds=time.perf_counter()-start,counters=c,hashes=hashes,solver={k:r[k] for k in ['acceptedSteps','rejectedSteps','nfev','njev','nlu','solverStarts','stateTransitions']},newtonIterations=d['newtonIterations'],newtonConvergenceFailures=d['newtonConvergenceFailures'])
    write(folder/'validation.json',record)
    print('RUN',label,'exact OK','skip',c['realSkips'],'original',c['nativeOriginalExecutions'],'target us',c['pathSeconds']/896*1e6,'Jac',d['jacobianSeconds'],flush=True)
    return record


def benchmark(pairs):
    run(False,False,'warm-control');run(False,True,'warm-skip')
    results=[]
    for i in range(pairs):
        for skip in ([False,True] if i%2==0 else [True,False]):results.append(run(False,skip,f'pair-{i+1}-'+('skip' if skip else 'control')))
    write(OUT/'performance.json',results)


if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('action',choices=['prepare','run','benchmark']);p.add_argument('--audit',action='store_true');p.add_argument('--skip',action='store_true');p.add_argument('--label');p.add_argument('--force',type=int,default=0);p.add_argument('--pairs',type=int,default=7);a=p.parse_args()
    if a.action=='prepare':prepare(a.audit,a.skip)
    elif a.action=='run':run(a.audit,a.skip,a.label,a.force)
    else:benchmark(a.pairs)