"""One fresh-process, CPU-only diagnosis of the stopped AO postprocessor.

Consumes the existing sealed first trajectory; never executes a neural trial,
changes the stopped protocol, or installs a delta. Diagnostic files are new.
"""
import argparse
import hashlib
import json
import os
from pathlib import Path
import sys
import time

ROOT = Path('/data/digital-drosophila')
EXP = ROOT / 'experiments/phase4ao-conditioning-interface'
MODEL = ROOT / 'models/male-cns-lif/olfactory-conditioning-lab-v1'
DIAG = ROOT / 'experiments/phase4ao-postprocess-isolation'
E_SHA = '757935a4815d9f625640df685184ed3fe44a8845323562796da7aa43cbb194e1'


def digest(p):
    h = hashlib.sha256()
    with Path(p).open('rb') as f:
        for data in iter(lambda: f.read(8 << 20), b''): h.update(data)
    return h.hexdigest()


def main(reg_sha, checkpoint_sha, authority_sha, stopped_sha):
    assert sys.dont_write_bytecode
    assert not any(k in sys.modules for k in ('runtime_light', 'runtime_continuation', 'phase4an_sealed_am_overlay'))
    regpath = EXP / 'protocol/preregistered.json'
    authority = EXP / 'receipts/D1--reference--epoch-000.authority.json'
    stopped = EXP / 'audit/stopped.json'
    checkpoint = EXP / 'runs/D1--reference--epoch-000/checkpoint-manifest.json'
    pins = {regpath: reg_sha, authority: authority_sha, stopped: stopped_sha, checkpoint: checkpoint_sha,
            MODEL / 'conditioning_epoch.py': E_SHA}
    assert all(digest(p) == h for p, h in pins.items())
    reg = json.loads(regpath.read_text())
    assert all(digest(p) == h for field in ('source_sha256', 'input_sha256', 'prepared_file_sha256') for p, h in reg[field].items())
    value = json.loads(authority.read_text())
    row = reg['plan']['calls'][0]
    assert row['row_id'] == 'D1--reference--epoch-000' and value['epoch_id'] == row['epoch_id']
    assert value['checkpoint_manifest_sha256'] == checkpoint_sha
    assert value['command'] == dict(steps=1200, compartment=reg['plan']['compartment_verbatim'], us_tick=900,
                                    d_command=1.0, frozen=False, trace_backend='torch')
    out = DIAG / 'postprocess'
    receipt = DIAG / 'diagnostic-receipt.json'
    assert DIAG.is_dir() and not out.exists() and not receipt.exists()
    os.environ['CUDA_VISIBLE_DEVICES'] = ''
    sys.path.insert(0, str(MODEL))
    import conditioning_epoch as E
    assert digest(E.__file__) == E_SHA
    started = time.perf_counter()
    chain = E.postprocess_epoch(authority, authority_sha, out)
    update = json.loads((out / 'update-receipt.json').read_text())
    assert all(digest(p) == h for p, h in pins.items())
    result = dict(status='PASS-ISOLATED-POSTPROCESS-DIAGNOSTIC-ONLY', registration_sha256=reg_sha,
        authority_sha256=authority_sha, checkpoint_manifest_sha256=checkpoint_sha, stopped_sha256=stopped_sha,
        diagnostic_script_sha256=digest(__file__), postprocess_source_sha256=E_SHA,
        result_files_sha256={str(p): digest(p) for p in out.iterdir()}, counts=update['counts'],
        before_overlay_sha256=chain['before_overlay_file_sha256'], after_overlay_sha256=chain['after_overlay_file_sha256'],
        wall_seconds=time.perf_counter() - started, CUDA_VISIBLE_DEVICES='',
        new_brain_runs=0, old_trajectory_rerun=False, weights_installed=False,
        original_protocol_remains_stopped=True, original_12_call_gate_passed=False,
        learning_success_claimed=False, candidate_generated=False,
        interpretation='The same pinned E succeeds in a fresh CPU process without the competing ordinary runtime loader cache. This is post-hoc diagnosis, not completion of AO.')
    with receipt.open('x', encoding='utf-8') as f:
        json.dump(result, f, sort_keys=True, indent=2, allow_nan=False); f.write('\n')
    print(json.dumps(dict(status=result['status'], receipt_sha256=digest(receipt), counts=result['counts'])), flush=True)


if __name__ == '__main__':
    p = argparse.ArgumentParser()
    for name in ('registration-sha256', 'checkpoint-sha256', 'authority-sha256', 'stopped-sha256'):
        p.add_argument('--' + name, required=True)
    a = p.parse_args()
    main(a.registration_sha256, a.checkpoint_sha256, a.authority_sha256, a.stopped_sha256)