"""R288-only typed replay cost-floor experiment; no production source changes."""
from pathlib import Path
import argparse,hashlib,json,os,re,shutil,subprocess,statistics,struct
import real_skip_experiment as base

ROOT=base.ROOT;HERE=Path(__file__).parent;OUT=ROOT/'test/r288-typed-replay-20260917'
OLD=base.OUT
original_core=base.core

def replace(s,a,b):return base.replace(s,a,b)
def body(s,name,change):return base.shadow.access.change_function(s,name,change)
def function_replace(s,name,new):
    a,b,e=base.ex.function_span(s,name);return s[:a]+new+s[e:]


def profile_core():
    s=original_core(False)
    pro='''
static uint64_t prof_ticks[24],prof_calls[24];
typedef struct {uint64_t start;int category;} Prof;
static void prof_end(Prof *s){prof_ticks[s->category]+=lp_tick()-s->start;prof_calls[s->category]++;}
#define PROF(c) Prof ps __attribute__((cleanup(prof_end)))={lp_tick(),c}
'''
    s=s.replace('static Plan plan;',pro+'\nstatic Plan plan;')
    for name,index in [('first_match',10),('map_slot',11),('scalar_lookup',12),('state_read_equal',13),('from_live',14)]:s=body(s,name,lambda b,i=index:'PROF('+str(i)+');'+b)
    s=s.replace('const Event *e=&plan.events[i];int slot=-1;','const Event *e=&plan.events[i];PROF(e->type);int slot=-1;')
    s=body(s,'replay_overlay',lambda b:replace('uint64_t header_start=lp_tick();'+b,'for(int i=0;i<SR_STATES;i++)mapping[i]=-1;','prof_ticks[17]+=lp_tick()-header_start;prof_calls[17]++;\nfor(int i=0;i<SR_STATES;i++)mapping[i]=-1;'))
    def split_copy(b):
        portions=[
            (18,'memset(dst,0,sizeof(*dst));memset(dst->states,0xa5,sizeof(dst->states));'),
            (19,'dst->context=*src;memcpy(dst->states,src->states,src->count*sizeof(*src->states));'),
            (20,'memcpy(dst->pipes,pipes,sizeof(dst->pipes));'),
            (21,'dst->memo=*src->jacobian;memcpy(dst->entries,src->jacobian->entries,src->jacobian->capacity*sizeof(*dst->entries));')]
        for tag,text in portions:b=replace(b,text,'{PROF('+str(tag)+');'+text+'}')
        start=b.index('dst->memo.entries=');b=b[:start]+'{PROF(22);'+b[start:]+'}'
        return b
    s=body(s,'from_live',split_copy)
    return s


def profile_runtime():
    s=(HERE/'real_skip_runtime.c').read_text()
    s=s.replace('static int force_every;','static int force_every;\nstatic FILE *plans_file;\nstatic uint64_t timer_ticks,dispatch_ticks;static volatile unsigned dispatch_sink;')
    s=body(s,'rr_start',lambda b:b+'\nplans_file=fopen("plans.bin","wb");if(!plans_file)fatal("plans");\nfor(int i=0;i<100000;i++){uint64_t t=lp_tick();timer_ticks+=lp_tick()-t;}\n')
    s=body(s,'rr_record_end',lambda b:b+'''
int saved=errno;SREnvironment env;save_environment(&env);
fwrite(&plan,1,sizeof(plan),plans_file);
uint64_t dispatch_start=lp_tick();unsigned checksum=0;
for(int repeat=0;repeat<100;repeat++)for(int i=0;i<plan.n;i++){
const Event *e=&plan.events[i];switch(e->type){
case QUERY:checksum+=(unsigned)e->slot;break;case SCALAR:checksum+=(unsigned)e->size;break;
case ALLOCATE:checksum++;break;case WRITE_STATE:case WRITE_PIPE:checksum+=(unsigned)e->offset;break;
case READ_STATE:case READ_PIPE:checksum+=(unsigned)e->size;break;case WRITE_OR:case VALID:checksum+=e->mask;break;default:checksum++;
}}
dispatch_ticks+=lp_tick()-dispatch_start;dispatch_sink=checksum;
restore_environment(&env);errno=saved;
''')
    s=body(s,'build_patch',lambda b:'PROF(15);'+b)
    s=body(s,'commit',lambda b:'PROF(16);'+b)
    s=body(s,'rr_finish',lambda b:b+'''
fclose(plans_file);FILE *pf=fopen("attribution.json","wb");
fprintf(pf,"{\\"frequency\\":%llu,\\"emptyTimerTicks\\":%llu,\\"emptyTimerCalls\\":100000,\\"dispatchTicks\\":%llu,\\"dispatchRepeats\\":100,\\"categories\\":[",(unsigned long long)frequency,(unsigned long long)timer_ticks,(unsigned long long)dispatch_ticks);
for(int i=0;i<24;i++)fprintf(pf,"%s{\\"id\\":%d,\\"ticks\\":%llu,\\"calls\\":%llu}",i?",":"",i,(unsigned long long)prof_ticks[i],(unsigned long long)prof_calls[i]);
fputs("]}",pf);fclose(pf);
''')
    return s


def minimal_sources(audit):
    # The audit build adds the minimal capture hooks to the certified access
    # source, allowing one physical baseline execution to feed BOTH schemas.
    directory=base.ACCESS if audit else base.BASE
    codes={n:(directory/n).read_text(encoding='utf-8') for n in ['properties.c','pipe.c']}
    s=codes['properties.c']
    if audit:
        # Native access hooks feed the compact recorder through ax_* wrappers.
        return codes
    s=body(s,'property_pt',lambda b:replace(replace(b,'return s;','{fp_query(cache,m,p,T,s);return s;}'),'return property_new(cache,m,p,T,scratch);','fp_query(cache,m,p,T,NULL);return property_new(cache,m,p,T,scratch);'))
    s=body(s,'property_new',lambda b:replace(b,'return s;','fp_allocate(cache,s,valid);return s;'))
    s=body(s,'native_temperature_ph_context',lambda b:'fp_unknown();'+b)
    s=body(s,'native_jacobian_scalar_get',lambda b:replace('fp_scalar_key(kind,medium_kind,inputs,count);'+b,'return 1;','fp_scalar_value(kind,medium_kind,inputs,count,*value);return 1;'))
    s=body(s,'native_jacobian_scalar_put',lambda b:'fp_scalar_value(kind,medium_kind,inputs,count,value);'+b)
    codes['properties.c']=s
    return codes


def fast_core(audit):
    if audit:
        s=original_core(True)
        s=s.replace('static SRContext overlay;','')
        s=body(s,'ax_query',lambda b:b+'\nif(phase==RECORD && strcmp(kind,"PT"))fp_unknown();')
        s=body(s,'ax_match',lambda b:b+'''
if(phase==RECORD && !strcmp(kind,"PT")){
Event *q=&plan.events[plan.query_index];double key[9];memcpy(key,q->data,sizeof(key));
NativeMedium m={q->offset,key[2],key[3],key[4],key[5],key[6],key[7],key[8]};fp_query(bound,&m,key[0],key[1],s);
}
''')
        s=body(s,'ax_new',lambda b:'if(phase==RECORD)fp_allocate(cache,s,valid);'+b)
        s=body(s,'ax_scalar',lambda b:b+'''
if(phase==RECORD){
 if(!strcmp(action,"get")){fp_scalar_key(kind,medium,keys,n);if(hit && value)fp_scalar_value(kind,medium,keys,n,*value);}
 else if(!strcmp(action,"put_attempt") && value)fp_scalar_value(kind,medium,keys,n,*value);
}
''')
        return s
    s=(HERE/'context_shadow_replay.c').read_text()
    s=s[:s.index('typedef struct {\n    int type,slot')]
    s=s.replace('static void save_environment','static __attribute__((unused)) void save_environment').replace('static void restore_environment','static __attribute__((unused)) void restore_environment')
    s+='''
static unsigned long long jac;
static NativePropertyCache *bound;
static NativePipeCache *bound_pipes;
static void fatal(const char *s){fprintf(stderr,"typed replay fatal: %s\\n",s);abort();}
'''
    return s


def fast_runtime(audit,direct=False):
    s=(HERE/'real_skip_runtime.c').read_text()
    s=s.replace('static int force_every;','static __attribute__((unused)) int force_every;\n#include "fast_replay.inc"')
    start=s.index('/* The translated write set');end=s.index('double rr_execute(',start)
    s=s[:start]+s[end:]
    execute=(HERE/'typed_replay_runtime.inc').read_text()
    s=function_replace(s,'rr_execute',execute)
    if audit:s=s.replace('void rr_start(void){','#include "typed_contract.inc"\n#include "typed_oracle.inc"\nvoid rr_start(void){')
    s=function_replace(s,'rr_jacobian','void rr_jacobian(void){jac++;fast.ready=0;\n#if RR_AUDIT\nplan.ready=0;\n#endif\n}')
    s=function_replace(s,'rr_record_begin','''void rr_record_begin(NativePropertyCache *p,NativePipeCache *pipes,const double *x){
baseline_start=lp_tick();bound=p;bound_pipes=pipes;fp_begin(p,pipes,x);
#if RR_AUDIT
ax_bind(p,pipes);ax_begin(jac,-1,16,0,x,4);
#endif
}''')
    s=function_replace(s,'rr_record_end','''void rr_record_end(double output){
fp_end(output);
#if RR_AUDIT
ax_end(&output,1);fp_contract_check();
#endif
baseline_ticks+=lp_tick()-baseline_start;
}''')
    if direct:
        s=s.replace('static __attribute__((unused)) int force_every;','static FILE *archived_plans;\nstatic __attribute__((unused)) int force_every;')
        s=body(s,'rr_start',lambda b:b+'\narchived_plans=fopen(getenv("RR_PLAN_ORACLE"),"rb");if(!archived_plans)fatal("archived plan oracle");\n')
        s=s.replace('ax_end(&output,1);fp_contract_check();','''
int saved_errno=errno;SREnvironment saved_env;save_environment(&saved_env);
if(fread(&plan,1,sizeof(plan),archived_plans)!=sizeof(plan))fatal("archived plan EOF");
uintptr_t previous_owner=plan.memo_binding;
for(int i=0;i<plan.n;i++){
Event *e=&plan.events[i];
if((e->type==READ_STATE || e->type==WRITE_STATE) && e->offset==(int)offsetof(NativePropertyState,jacobian)){
uintptr_t pointer;memcpy(&pointer,e->data,sizeof(pointer));if(pointer!=previous_owner)fatal("archived memo binding");memcpy(e->data,&fast.owner,sizeof(fast.owner));
}}
plan.memo_binding=(uintptr_t)fast.owner;phase=OFF;
restore_environment(&saved_env);errno=saved_errno;fp_contract_check();
''')
        s=body(s,'rr_finish',lambda b:b+'\nif(fgetc(archived_plans)!=EOF){fatal("remaining archived plans");}\nfclose(archived_plans);\n')
    s=s.replace('(unsigned long long)sizeof(plan),(unsigned long long)sizeof(overlay),(unsigned long long)sizeof(patches)', '(unsigned long long)sizeof(fast),(unsigned long long)sizeof(Pending),(unsigned long long)sizeof(Pending)')
    s=body(s,'rr_finish',lambda b:b+'''
FILE *fextra=fopen("typed-summary.json","wb");
fprintf(fextra,"{\\"metadataBytes\\":%llu,\\"captureScratchBytes\\":%llu,\\"pendingBytes\\":%llu,\\"contractChecks\\":%llu,\\"oracleChecks\\":%llu,\\"negativeChecks\\":%llu}",(unsigned long long)sizeof(fast),(unsigned long long)sizeof(capture),(unsigned long long)sizeof(Pending),
#if RR_AUDIT
contract_checks,oracle_checks,negative_checks
#else
0ULL,0ULL,0ULL
#endif
);fclose(fextra);
''')
    return s


def contract_source():
    data=(OUT/'P/perf-skip-run/plans.bin').read_bytes();schemas={}
    for off in range(0,len(data),90224):
        count=struct.unpack_from('<Q',data,off+90112+56)[0];n=struct.unpack_from('<i',data,off+90112+68)[0];path=struct.unpack_from('<i',data,off+90112+88)[0]
        entries=[]
        for i in range(n):
            t,s,o,z,a,m=struct.unpack_from('<iiiiiI',data,off+i*176)
            if t in (1,2,3,7,8):s-=count
            entries.append((t,s,o,z,a,m))
        if path in schemas:assert schemas[path]==entries
        else:schemas[path]=entries
    assert sorted(schemas)==[0,1] and len(schemas[0])==12 and len(schemas[1])==132
    prefix='typedef struct {int type,slot,offset,size,aux;unsigned mask;} Contract;\n'
    for path,events in schemas.items():prefix+='static const Contract schema'+str(path)+'[]={'+','.join('{'+','.join(map(str,e))+'}' for e in events)+'};\n'
    return prefix+(HERE/'typed_contract_check.inc').read_text()


def prepare(mode,audit=False):
    certified=json.loads((OLD/'perf-skip/build.json').read_text())['originalHashes']
    for name,expected in certified.items():assert hashlib.sha256((base.BASE/name).read_text(encoding='utf-8').encode()).hexdigest()==expected,name
    root=OUT/mode;root.mkdir(parents=True,exist_ok=True)
    staging=root/'source';staging.mkdir(exist_ok=True)
    for name in ['context_shadow_replay.c','context_shadow_replay.h','real_skip.h']:shutil.copyfile(HERE/name,staging/name)
    runtime=fast_runtime(audit,mode=='D') if mode in 'CD' else profile_runtime() if mode=='P' else (HERE/'real_skip_runtime.c').read_text()
    (staging/'real_skip_runtime.c').write_text(runtime,encoding='utf-8')
    if mode in 'CD':
        header=(staging/'real_skip.h').read_text();header+='\n'+(HERE/'typed_replay_hooks.h').read_text();(staging/'real_skip.h').write_text(header)
        work=root/('audit-skip' if audit else 'perf-skip');work.mkdir(exist_ok=True)
        shutil.copyfile(HERE/'typed_replay_core.inc',work/'fast_replay.inc')
        if audit:
            (work/'typed_contract.inc').write_text(contract_source(),encoding='utf-8')
            shutil.copyfile(HERE/'typed_oracle.inc',work/'typed_oracle.inc')
        access=staging/'capture';access.mkdir(exist_ok=True)
        for n in ['context_access_diag.h','context_fallback_diag.h']:shutil.copyfile(base.ACCESS/n,access/n)
        for n,code in minimal_sources(audit and mode!='D').items():(access/n).write_text(code,encoding='utf-8')
    base.OUT=root;base.HERE=staging
    if mode in 'CD':base.ACCESS=access
    base.core=lambda enabled:fast_core(enabled) if mode in 'CD' else profile_core() if mode=='P' else original_core(enabled)
    if not audit:
        gate=root/'audit-skip-run';gate.mkdir(exist_ok=True)
        if mode!='C':shutil.copyfile(OLD/'audit-skip-run/validation.json',gate/'validation.json')
        else:assert json.loads((gate/'validation.json').read_text())['exact']
    # original_core reads base.HERE; the source copies are byte-identical.
    base.prepare(audit,mode!='A')
    if mode in 'CD':
        manifest=root/('audit-skip' if audit else 'perf-skip')/'build.json';info=json.loads(manifest.read_text())
        info['guardMatchesShadow']=False;info['wholeContextGuardUnchanged']=True
        info['typedGuardValidation']='See C/D full-trajectory contracts, generic oracle and forced-reject audits.'
        base.write(manifest,info)
    if mode=='C' and not audit:
        cc,_,_=base.ex.builder.toolchain();nm=Path(cc).with_name('nm.exe' if os.name=='nt' else 'nm')
        symbols=subprocess.check_output([str(nm),'--defined-only',str(root/'perf-skip/model.exe')],text=True)
        assert not re.search(r'(?m)\b(?:plan|overlay|before_reject|after_reject|oracle_generic|oracle_typed|replay_overlay|ax_access)$',symbols)
        base.write(root/'perf-skip/stripped-audit-proof.json',dict(noGenericInterpreter=True,noFullOverlay=True,noAuditOracle=True))


def run(mode,audit=False,label=None,force=0):
    base.OUT=OUT/mode
    if mode=='D':os.environ['RR_PLAN_ORACLE']=str(OUT/'P/perf-skip-run/plans.bin')
    if audit:
        expected=base.OUT/'audit-control-run';expected.mkdir(exist_ok=True)
        target=expected/'audit.bin'
        if not target.exists():os.link(OLD/'audit-control-run/audit.bin',target)
    result=base.run(audit,mode!='A',label,force)
    result['mode']=mode
    if mode in 'CD':
        extra=json.loads((base.OUT/result['label']/'typed-summary.json').read_text())
        assert extra['metadataBytes']==240 and extra['captureScratchBytes']==96 and extra['pendingBytes']==456
        assert extra['contractChecks']==extra['oracleChecks']==(896 if audit else 0)
        assert extra['negativeChecks']==(20 if audit and not force else 0)
        result['typed']=extra
    base.write(base.OUT/result['label']/'validation.json',result)
    return result


def benchmark(rounds):
    for mode in 'CD':
        for label in ['audit-skip-run','forced-reject']:assert json.loads((OUT/mode/label/'validation.json').read_text())['exact']
    rows=[]
    for mode in 'ABC':run(mode,label='warm')
    for i in range(rounds):
        order='ABC'[i%3:]+'ABC'[:i%3]
        for mode in order:rows.append(run(mode,label=f'round-{i+1}'))
        base.write(OUT/'performance.json',rows)


if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('action',choices=['prepare','run','benchmark']);p.add_argument('--mode',choices=list('ABCPD'),default='C');p.add_argument('--audit',action='store_true');p.add_argument('--label');p.add_argument('--force',type=int,default=0);p.add_argument('--rounds',type=int,default=9);a=p.parse_args()
    if a.action=='prepare':prepare(a.mode,a.audit)
    elif a.action=='run':run(a.mode,a.audit,a.label,a.force)
    else:benchmark(a.rounds)