"""position379 only: diagnose the actual numerical tail before any real skip."""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import argparse,hashlib,json,os,shutil,subprocess,time
import local_probe_experiment as ex

ROOT=ex.ROOT
HERE=Path(__file__).parent
BASE=ROOT/'test/local-probe-20260917/worker'
OUT=ROOT/'test/position379-tail-20260917'

def replace(s,a,b):
    assert s.count(a)==1,(a,s.count(a))
    return s.replace(a,b)

def prepare(trace=True,tail_only=False):
    name='tail-only-trace' if tail_only else 'diag-trace' if trace else 'diag-perf'
    work=OUT/name;work.mkdir(parents=True,exist_ok=True)
    src={p.name:p.read_text(encoding='utf-8') for p in BASE.glob('*.c')}
    hashes={k:hashlib.sha256(v.encode()).hexdigest() for k,v in src.items()}
    plan=json.loads((ROOT/'test/context-fallback-20260917/plan.json').read_text(encoding='utf-8'))
    s=src['model.c'];a,b,e=ex.function_span(s,'model_eval_local_internal');body=s[a:e];op=plan['code'][379][0]
    assert body.count(op)==2
    body=body.replace(op,'kd_scope=1;'+op+'kd_scope=0;')
    body=replace(body,'return 1;}','kd_eval_exit(t,y,dy,w,1,properties,pipe_cache,jacobian);return 1;}')
    src['model.c']=s[:a]+body+s[e:]
    s=src['properties.c'];a,b,e=ex.function_span(s,'state_valve');body=s[a:e]
    boundary='    double r=fmax(pd/p,0),critical=pow(2*g/(g+1),1/(1-g)),eff;'
    body=replace(body,boundary,'''    KdRecord kd_record;uint64_t kd_t=0;
    if(kd_scope){uint64_t kd_lookup_start=kd_clock();kd_before(&kd_record,p,T,pd,g,rho);kd_record.lookup=kd_clock()-kd_lookup_start;
        kd_t=kd_clock();
    }
'''+boundary)
    end=body.rfind('}')
    body=body[:end]+'    if(kd_scope){uint64_t kd_end=kd_clock();kd_after(&kd_record,*cm,*vel,kd_end-kd_t,1);}\n'+body[end:]
    if tail_only:
        # Cross-check K without key construction, lookup, or environment reads
        # immediately before the timed numerical work. Post-tail records in this
        # mode MUST NOT be used for memo/FP eligibility or G/H estimates.
        before='uint64_t kd_lookup_start=kd_clock();kd_before(&kd_record,p,T,pd,g,rho);kd_record.lookup=kd_clock()-kd_lookup_start;'
        body=replace(body,before,'')
        body=replace(body,'uint64_t kd_end=kd_clock();kd_after(', 'uint64_t kd_end=kd_clock();kd_before(&kd_record,p,T,pd,g,rho);kd_record.lookup=0;kd_after(')
    src['properties.c']=s[:a]+body+s[e:]
    src['common.c']=replace(replace(src['common.c'],'lp_start();','lp_start();kd_start();'),'lp_finish();','lp_finish();kd_finish();')
    src['cvode_solver.c']=replace(src['cvode_solver.c'],'{lp_color=-1;uint64_t start=lp_tick();','{lp_color=-1;kd_new_jac();uint64_t start=lp_tick();')
    src['cvode_solver.c']=replace(src['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));kd_matrix(t,N_VGetArrayPointer(y),SUNDenseMatrix_Data(matrix));}')
    src['valve_tail_diag.c']=(HERE/'valve_tail_diag.c').read_text(encoding='utf-8')
    for name in ('model.h','local_probe.h'):shutil.copyfile(BASE/name,work/name)
    shutil.copyfile(HERE/'valve_tail_diag.h',work/'valve_tail_diag.h')
    cc,sun,_=ex.builder.toolchain();flags,libs,dlls,exe=ex.builder.platform_build_inputs(sun)
    flags+=['-DLP_OBSERVE=0',f'-DKD_TRACE={int(trace)}']
    def compile_one(item):
        name,code=item;path=work/name;path.write_text('#include "valve_tail_diag.h"\n'+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,src.items()))
    log=[];ex.builder._command([cc,*flags,*[str(o) for o,_ 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([l for _,l in objects],[])+log),encoding='utf-8')
    assert src['local_probe_support.c']==(BASE/'local_probe_support.c').read_text(encoding='utf-8')
    ex.write(work/'build.json',dict(originalHashes=hashes,skip=False,trace=trace,tailOnlyTiming=tail_only,wholeContextGuardUnchanged=True,sourceHashes={p.name:hashlib.sha256(p.read_bytes()).hexdigest() for p in work.glob('*.c')}))
    print('BUILT',work.name,flush=True)

def run(label,worker='diag-trace',matrices=False,disabled=False):
    folder=OUT/label;folder.mkdir(parents=True,exist_ok=True);env=os.environ.copy();env['LOCAL_PROBE_MASK']='0x7ffffff'
    if matrices:env['KD_MATRICES']='1'
    if disabled:env['KD_DISABLED']='1'
    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:r=subprocess.run(args,cwd=folder,env=env,stdout=subprocess.PIPE,stderr=f,timeout=330,creationflags=subprocess.CREATE_NO_WINDOW)
    assert r.returncode==0,(label,r.returncode,(folder/'stderr.log').read_text(encoding='utf-8')[-4000:])
    result=json.loads((folder/'result.json').read_text(encoding='utf-8'));d=json.loads((folder/'probe.json').read_text(encoding='utf-8'))
    reference=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(result[k]==reference[k] for k in keys),[(k,result[k],reference[k]) for k in keys if result[k]!=reference[k]]
    for k in ['newtonIterations','newtonConvergenceFailures','modelCalls','groups','contextCopiedBytes','contextComparedBytes']:assert d[k]==reference['diagnostic'][k],k
    hashes={}
    for name in ['states','outputs','events']+(['jacobians'] if matrices else []):
        hashes[name]=hashlib.sha256((folder/(name+'.bin')).read_bytes()).hexdigest()
        expected=hashlib.sha256((BASE.parent/'all-audit/jacobians.bin').read_bytes()).hexdigest() if name=='jacobians' else reference[name+'Sha256']
        assert hashes[name]==expected,name
    data=json.loads((folder/'tail.json').read_text());assert data['jacobians']==896
    record=dict(label=label,worker=worker,exact=True,hashes=hashes,counters={k:result[k] for k in keys},newtonIterations=d['newtonIterations'],newtonConvergenceFailures=d['newtonConvergenceFailures'],tail=data,wall=result['solveSeconds'],cpu=result['solveCpuSeconds'],jacobian=d.get('jacobianSeconds'),processSeconds=time.perf_counter()-start)
    ex.write(folder/'validation.json',record);print(json.dumps(dict(label=label,exact=True,tail=data,wall=record['wall'],cpu=record['cpu']),ensure_ascii=True),flush=True)
    return record

if __name__=='__main__':
    p=argparse.ArgumentParser();p.add_argument('action',choices=['prepare','run']);p.add_argument('--label',default='diagnostic-0');p.add_argument('--worker',default='diag-trace');p.add_argument('--no-trace',action='store_true');p.add_argument('--tail-only',action='store_true');p.add_argument('--matrices',action='store_true');p.add_argument('--disabled',action='store_true');a=p.parse_args()
    if a.action=='prepare':prepare(not a.no_trace,a.tail_only)
    else:run(a.label,a.worker,a.matrices,a.disabled)