"""Fresh CPU-process revalidation of a sealed epoch's actual update.

This does not run a brain or E.postprocess_epoch and writes only one exclusive
receipt. The caller must pin the source, authority and successful process to its
registration. A supplied SHA is not independent scientific authorization.
"""
from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
import os
from pathlib import Path
import sys

import numpy as np

LAB = Path(__file__).resolve().parent.parent/'olfactory-conditioning-lab-v1'
E_SHA = '757935a4815d9f625640df685184ed3fe44a8845323562796da7aa43cbb194e1'
Q_SHA = 'b88fd9a88d9f92d5c4d042b90f2129b015e8d4bd810c41351ee573241e2d4c2e'
U_SHA = '96efff678cff910654eb08e86f1a17789e4e5589e340345f419c7027a7d547e2'
CMP_SHA = '5a2403f11bc08d65a9d242ac2cf9686b850e0c3680d65659f785e9b6b50c0238'
FILES = ('chain-receipt.json','update-receipt.json','delta-after.npz','overlay-manifest.json')


def require(ok,message):
    if not ok: raise ValueError(message)


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


def same(a,b): return a.dtype==b.dtype and a.shape==b.shape and a.tobytes()==b.tobytes()


def unlinked(path):
    path=Path(path).absolute()
    require(not any(p.is_symlink() or (hasattr(p,'is_junction') and p.is_junction())
                    for p in (path,*path.parents)), 'Linked receipt/update path forbidden')
    return path


def load(name,path,expected):
    raw=Path(path).read_bytes()
    require(hashlib.sha256(raw).hexdigest()==expected and name not in sys.modules,'Unverified source/module binding')
    spec=importlib.util.spec_from_file_location(name,path);module=importlib.util.module_from_spec(spec)
    sys.modules[name]=module
    exec(compile(raw,str(path),'exec'),module.__dict__)
    return module


def chain_binding(chain,authority,authority_sha,pins,paths,checkpoint,hashes):
    require(chain.get('schema')=='phase4an-epoch-postprocess-chain-v1' and chain.get('complete') is True,
            'Incomplete epoch update chain')
    for flag in ('weights_installed','brain_executed','learning_pass_claimed','candidate_generated'):
        require(chain.get(flag) is False,'Wrong chain scope: '+flag)
    require(chain['coordinator_source_sha256']==E_SHA and chain['authority_sha256']==authority_sha and
        chain['epoch_id']==authority['epoch_id'] and chain['command']==authority['command'] and
        chain['diagnostic_mode']==authority['diagnostic_mode'],'Epoch authority/command differs')
    require(chain['source_neural_checkpoint_dir']==str(checkpoint) and
        chain['source_neural_checkpoint_manifest_sha256']==authority['checkpoint_manifest_sha256'] and
        chain['input_files_sha256']=={str(p):v for p,v in pins.items()},'Actual checkpoint/input binding differs')
    for field,role in (('source_state_file_sha256','state'),('source_spec_sha256','spec'),
                       ('source_requests_sha256','requests'),('before_overlay_file_sha256','overlay_before'),
                       ('before_overlay_manifest_sha256','overlay_manifest_before')):
        require(chain[field]==pins[paths[role]],'Chain input differs: '+role)
    expected={name:hashes[name] for name in FILES if name!='chain-receipt.json'}
    require(chain['output_files_sha256']==expected and chain['after_overlay_file_sha256']==hashes['delta-after.npz'] and
            chain['after_overlay_manifest_sha256']==hashes['overlay-manifest.json'],'Chain output differs')


def validate(epoch_authority,authority_sha,update_dir,chain_sha,receipt_file):
    require(sys.dont_write_bytecode and 'torch' not in sys.modules and 'runtime_light' not in sys.modules,
            'Run in a fresh Python -B CPU process')
    os.environ['CUDA_VISIBLE_DEVICES']=''
    e=load('_conditioning_execution_E_validator',LAB/'conditioning_epoch.py',E_SHA)
    load('engineering_compare',LAB/'engineering_compare.py',CMP_SHA)
    q=load('_conditioning_execution_Q_validator',LAB/'conditioning_preflight_compare.py',Q_SHA)
    value,paths,pins,checkpoint,authority_path=e._authority(epoch_authority,authority_sha)
    c,r,u=e._production_dependencies(paths,pins)  # The only R is C.R; no ordinary runtime loader.
    directory=unlinked(update_dir).resolve(strict=True);out=unlinked(receipt_file)
    protected={checkpoint,directory,*[p.parent for p in pins]}
    require(out.parent.is_dir() and not out.exists() and not any(out.resolve().is_relative_to(p) for p in protected),
            'Receipt must be new and outside every frozen input/output directory')
    require(set(p.name for p in directory.iterdir())==set(FILES) and
            all((directory/n).is_file() and not (directory/n).is_symlink() for n in FILES),'Expected exact four update files')
    hashes={n:digest(directory/n) for n in FILES}
    require(e._sha(chain_sha) and hashes['chain-receipt.json']==chain_sha,'Chain differs from external SHA')
    own=digest(__file__);chain=e._read(directory/'chain-receipt.json')
    chain_binding(chain,value,authority_sha,pins,paths,checkpoint,hashes)
    parent=c.load_checkpoint(checkpoint,value['checkpoint_manifest_sha256'])
    command=value['command'];spec=e._read(paths['spec'])
    r.validate_spec(spec,paths['spec'].parent)
    require(parent['fresh'] is True and parent['global_tick_offset']==0 and parent['steps']==command['steps']==spec['steps'] and
            parent['summary']['diagnostic_mode']==value['diagnostic_mode'],'Wrong actual epoch/state scope')
    for role in ('spec','requests','overlay_before','overlay_manifest_before'):
        require(parent['provenance']['input_sha256'].get(str(paths[role]))==pins[paths[role]],'Runtime input differs: '+role)
    require(parent['manifest']['files_sha256']['state.npz']==pins[paths['state']],'State not sealed by this checkpoint')
    _,before,before_meta,_=r.load_overlay(paths['overlay_before'],paths['overlay_manifest_before'])
    _,after,after_meta,_=r.load_overlay(directory/'delta-after.npz',directory/'overlay-manifest.json')
    graph=Path(spec['graph_dir']).resolve();apl=Path(spec['apl_bundle_dir']).resolve()
    for location in [graph/n for n in e.GRAPH_FILES]+[apl/n for n in e.APL_FILES]+[Path(spec['coefficients_path']).resolve()]:
        require(location in pins,'Missing complete graph/APL/coefficient pin')
    require(paths['canonical_kc']==apl/'kc_globals.npy','Canonical KC path differs')
    for role,prefix in (('mask','mask'),('base_candidate','base_candidate'),('am_overlay_source','am_overlay')):
        require(Path(before_meta[prefix+'_path']).resolve()==paths[role] and before_meta[prefix+'_sha256']==pins[paths[role]],'Overlay input differs')
    aj=Path(before_meta['aj_model_dir']).resolve()
    for relative,sha in r.PINS.items():require(pins.get((aj.parent/relative).resolve())==sha,'Missing original dependency pin')
    mapping=e._npz(paths['candidate_map'],('edge_id','source_neuron_idx','target_neuron_idx','target_compartment_verbatim'))
    mask=np.load(paths['mask'],allow_pickle=False);kc=np.load(paths['kc'],allow_pickle=False)
    canonical=np.load(paths['canonical_kc'],allow_pickle=False);base=np.load(paths['base_candidate'],allow_pickle=False)
    e._vector(mask,'int64',14232,'mask');e._vector(kc,'int32',4064,'KC');e._vector(canonical,'int32',4064,'canonical KC')
    e._vector(base,'float32',14232,'base')
    require(np.all((mask>=0)&(mask<25582938)) and np.all(mask[1:]>mask[:-1]) and
            np.all((kc>=0)&(kc<166700)) and len(np.unique(kc))==len(kc),
            'Candidate/KC range or uniqueness differs')
    require(same(kc,canonical) and same(mask,mapping['edge_id']) and same(mask,before['edge_id']) and same(mask,after['edge_id']),
            'KC/candidate membership differs')
    require(same(base,before['base_candidate_mv']) and same(base,after['base_candidate_mv']),'Original base changed')
    maps={}
    try:
        for n,dtype,count in (('nodes','int32',166700),('edge_id','int64',25582938),('src','int32',25582938),('dst','int32',25582938)):
            maps[n]=np.load(graph/(n+'.npy'),allow_pickle=False,mmap_mode='r');e._vector(maps[n],dtype,count,n)
        require(np.array_equal(maps['nodes'],np.arange(166700,dtype=np.int32)) and same(maps['edge_id'][mask],mask),'Not canonical full graph')
        require(same(mapping['source_neuron_idx'],maps['src'][mask]) and same(mapping['target_neuron_idx'],maps['dst'][mask]),
                'Candidate endpoints differ from actual connectome')
    finally:
        for a in maps.values():a._mmap.close()
    state=e._npz(paths['state'],('step','neuron_idx','counts','v','g','ref'))
    require(all(same(state[n],parent['state'][n]) for n in state),'Actual sealed event/state mismatch')
    require(parent['overlay']['delta_raw_sha256']==r._raw(before['plastic_delta_mv']) and
            chain['runtime_delta_before_raw_sha256']==r._raw(before['plastic_delta_mv']) and
            chain['runtime_effective_weights_raw_sha256']==parent['summary']['effective_weights_raw_sha256'], 'Actual before installation differs')
    update=e._read(directory/'update-receipt.json')
    require(update['payload_sha256']==chain['update_receipt_payload_sha256'],'Update receipt chain differs')
    u.verify_receipt(update,step=state['step'],neuron_idx=state['neuron_idx'],kc=kc,candidate_edge_id=mask,
        candidate_sources=mapping['source_neuron_idx'],candidate_targets=mapping['target_neuron_idx'],
        candidate_compartments=mapping['target_compartment_verbatim'],delta_before=before['plastic_delta_mv'],base_candidate=base,
        steps=command['steps'],n_neurons=166700,compartment=command['compartment'],us_tick=command['us_tick'],
        d_command=command['d_command'],frozen=command['frozen'],trace_backend=command['trace_backend'],
        bridge_source=paths['eligibility_bridge_source'],phase4j_source=paths['phase4j_source'])
    binding=q.audit_delta_binding(before['plastic_delta_mv'],after['plastic_delta_mv'],after['plastic_delta_mv'],
        update_receipt=update,parent_snapshot_sha256=pins[paths['overlay_before']],child_manifest=after_meta)
    require(binding['pass_all'],'Actual after array differs from recomputed receipt')
    import torch
    require(os.environ['CUDA_VISIBLE_DEVICES']=='' and not torch.cuda.is_initialized(),'CPU verifier initialized CUDA')
    e._check_files(pins);c.load_checkpoint(checkpoint,value['checkpoint_manifest_sha256'])
    for file,sha in (('conditioning_epoch.py',E_SHA),('conditioning_preflight_compare.py',Q_SHA),('engineering_compare.py',CMP_SHA)):
        require(digest(LAB/file)==sha,'Verifier dependency changed')
    require(hashes=={n:digest(directory/n) for n in FILES} and own==digest(__file__),'Source/update changed during validation')
    receipt=dict(schema='conditioning-update-cpu-validation-v1',complete=True,epoch_authority_sha256=authority_sha,
        parent_checkpoint_manifest_sha256=value['checkpoint_manifest_sha256'],update_files_sha256=hashes,
        before_overlay_sha256=pins[paths['overlay_before']],before_manifest_sha256=pins[paths['overlay_manifest_before']],
        after_overlay_sha256=hashes['delta-after.npz'],after_manifest_sha256=hashes['overlay-manifest.json'],
        before_delta_raw_sha256=r._raw(before['plastic_delta_mv']),after_delta_raw_sha256=r._raw(after['plastic_delta_mv']),
        update_receipt_payload_sha256=update['payload_sha256'],validator_source=dict(path=str(Path(__file__).resolve()),sha256=own),
        frozen_U_source_sha256=U_SHA,frozen_U_verify_receipt_recomputed=True,actual_sealed_events_used=True,
        full_candidate_mapping_verified=True,brain_executed=False,learning_pass_claimed=False)
    e._save(out,receipt)
    return receipt


if __name__=='__main__':
    parser=argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--authority',required=True);parser.add_argument('--authority-sha256',required=True)
    parser.add_argument('--update-dir',required=True);parser.add_argument('--chain-sha256',required=True)
    parser.add_argument('--receipt',required=True);args=parser.parse_args()
    validate(args.authority,args.authority_sha256,args.update_dir,args.chain_sha256,args.receipt)
    print(json.dumps(dict(complete=True,receipt=args.receipt,sha256=digest(args.receipt))),flush=True)