"""Independent, opt-in access recorder for positions 16 and 52 only.

Uses the existing fallback diagnostic worker as its source; never edits the
production kernels or changes the local-probe guard/replay policy.
"""
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
import argparse, hashlib, json, re, shutil, time
import diagnose_context_fallback as dx

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


def instrument(body):
    """Wrap reviewed lvalues first, then actual rvalues; keep short circuits."""
    saved = []
    def hold(text):
        saved.append(text)
        return f'AXHOLD{len(saved)-1}ZZ'
    field = r'(?:s|up|down|cache|m|a|b)->\w+'
    def reads(text):
        # An address expression is not a value read (e.g. &s->rho).
        text = re.sub(r'(?<![\w&])' + field, lambda m: 'AX_R(' + m[0] + ')', text)
        # Taking &states[i] still consumes the states base pointer.
        text = text.replace('&cache->states[', '&AX_R(cache->states)[')
        return 'AX_R(*m)' if text.strip() == '*m' else text
    body = re.sub(r'([su]\w*->valid)&(NATIVE_PROPERTY_\w+)',
                  lambda m: hold(f'ax_test({m[1].split("->")[0]},{m[2]},__func__)'), body)
    body = body.replace('cache->count++', hold('AX_INC(cache->count)'))
    # Struct reset is an explicit write, including equal-valued fields.
    body = body.replace('*s=(NativePropertyState){0};', hold('AX_W(*s,((NativePropertyState){0}));'))
    def write(m):
        lhs, operator, rhs = m.groups()
        return hold(('AX_OR' if operator == '|=' else 'AX_W') + f'({lhs},{reads(rhs)});')
    body = re.sub(r'(' + field + r')\s*(\|=|=(?!=))\s*([^;]+);', write, body)
    body = reads(body)
    for i, value in enumerate(saved):
        body = body.replace(f'AXHOLD{i}ZZ', value)
    # Fail closed if a field write evaded the transformation.
    assert not re.search(field + r'\s*(?:\|=|=(?!=)|\+\+)', body), body
    return body


def change_function(source, name, transform):
    a, b, e = dx.ex.function_span(source, name)
    return source[:b+1] + transform(source[b+1:e-1]) + source[e-1:]


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, 'First prepare the existing context-fallback diagnostic worker.'
    before = {k: hashlib.sha256(v.encode()).hexdigest() for k, v in sources.items()}
    s = sources['properties.c']
    def pt(body):
        body = 'ax_query("PT",m,p,T);\n' + body
        body = replace(body, 'same_medium(&s->medium,m))return s;', 'same_medium(&s->medium,m)){ax_match("PT",s);return s;}')
        body = replace(body, 'return property_new(cache,m,p,T,scratch);', 'ax_match("PT",NULL);return property_new(cache,m,p,T,scratch);')
        return body
    s = change_function(s, 'property_pt', pt)
    def ph(body):
        body = 'ax_query("PH",m,p,h);\n' + body
        body = replace(body, 'observe_temperature(cache->temperatures,m,s->T,3);return s->T;', 'ax_match("PH",s);observe_temperature(cache->temperatures,m,s->T,3);return s->T;')
        return replace(body, 'double T;', 'ax_match("PH",NULL);double T;')
    s = change_function(s, 'native_temperature_ph_context', ph)
    s = change_function(s, 'property_new', lambda body: replace(body, '*s=(NativePropertyState){0};', 'ax_new(cache,s,valid);*s=(NativePropertyState){0};'))
    functions = ['same_medium', 'property_new', 'property_pt', 'property_density', 'property_viscosity',
                 'native_temperature_ph_context', 'local_isentropic', 'isentropic', 'state_valve',
                 'native_density', 'native_temperature_ph', 'native_viscosity']
    for name in functions:
        s = change_function(s, name, instrument)
    # Scalar get can write directly into s->rho through an output pointer.
    s = change_function(s, 'native_jacobian_scalar_get', lambda body: replace(
        replace(body, '*value=entry->value;', 'AX_W(*value,entry->value);ax_scalar("get",cache,kind,medium_kind,inputs,count,1,value);'),
        'return 0;', 'ax_scalar("get",cache,kind,medium_kind,inputs,count,0,NULL);return 0;'))
    s = change_function(s, 'native_jacobian_scalar_put', lambda body:
        'ax_scalar("put_attempt",cache,kind,medium_kind,inputs,count,-1,&value);\n' + body)
    sources['properties.c'] = s
    for name in ['native_pipe_flow_context', 'native_pipe_flow_cached_context']:
        sources['pipe.c'] = change_function(sources['pipe.c'], name, instrument)
    s = sources['local_probe_support.c']
    s = replace(s, 'dx_region=-1;dx_boundary(0,-2,p,pipes);', 'ax_bind(p,pipes);dx_region=-1;dx_boundary(0,-2,p,pipes);')
    s = replace(s, 'void dx_op_begin(int pos,const double *inputs){',
                'void dx_op_begin(int pos,const double *inputs){ax_begin(dx_jac-1,lp_color,pos,dx_time,inputs,dx_in_offset[pos+1]-dx_in_offset[pos]);')
    s = replace(s, 'void dx_op_end(int pos,const double *outputs,NativePropertyCache *p,NativePipeCache *pipes){',
                'void dx_op_end(int pos,const double *outputs,NativePropertyCache *p,NativePipeCache *pipes){ax_end(outputs,dx_out_offset[pos+1]-dx_out_offset[pos]);')
    s = replace(s, 'void dx_finish(void){', 'void dx_finish(void){ax_finish();')
    sources['local_probe_support.c'] = s
    sources['model.c'] = replace(sources['model.c'], 'lp_active=-1;dx_eval_end();return result;', 'lp_active=-1;ax_result(result);dx_eval_end();return result;')
    sources['context_access_diag.c'] = (HERE/'context_access_diag.c').read_text(encoding='utf-8')
    for path in BASE.glob('*.h'):
        shutil.copyfile(path, work/path.name)
    shutil.copyfile(HERE/'context_access_diag.h', work/'context_access_diag.h')
    cc, sun, _ = dx.ex.builder.toolchain()
    flags, libs, dlls, exe = 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
        if '#define _WIN32_WINNT 0x0600' in code:
            code = code.replace('#define _WIN32_WINNT 0x0600', '#define _WIN32_WINNT 0x0600\n#include "context_access_diag.h"', 1)
        else:
            code = '#include "context_access_diag.h"\n' + code
        path.write_text(code, encoding='utf-8', newline='\n')
        obj = path.with_suffix('.o'); log = []
        dx.ex.builder._command([cc, *flags, '-I', str(work), '-I', str(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 = []
    dx.ex.builder._command([cc, *flags, *[str(o) for o, _ in objects], *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')
    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, functions=functions))
    print('BUILT access worker', flush=True)


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('action', choices=['prepare', 'run'])
    args = parser.parse_args()
    if args.action == 'prepare':
        prepare()
    else:
        dx.OUT = OUT
        dx.run('audit', mode=1, matrices=True)