import hashlib
import importlib.util
import json
from pathlib import Path

ROOT=Path('/data/digital-drosophila')
MODEL=Path(__file__).resolve().parent
EXP=ROOT/'experiments/phase4ag-author-apl-r2'
AF_MODEL=ROOT/'models/male-cns-lif/author-fullgraph-lab-v1'
AF_EXP=ROOT/'experiments/phase4af-author-fullgraph'
AD_EXP=ROOT/'experiments/phase4ad-author-semantics'
CFG=json.loads((MODEL/'config.json').read_text(encoding='utf-8'))

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

def save(path,value):
    path=Path(path);path.parent.mkdir(parents=True,exist_ok=True)
    with path.open('x',encoding='utf-8') as f:json.dump(value,f,indent=2,ensure_ascii=False,allow_nan=False);f.write('\n')

def sources():
    return {p.name:digest(p) for p in sorted(MODEL.iterdir()) if p.is_file() and p.name not in ('manifest.json','checksums.sha256')}

def integrity():
    spec=importlib.util.spec_from_file_location('frozen_af_context',AF_MODEL/'context.py')
    mod=importlib.util.module_from_spec(spec);spec.loader.exec_module(mod)
    result=mod.integrity()
    import sys
    sys.path.insert(0,str(ROOT/'models/male-cns-lif/terminal-scope-diagnostic-v1'))
    from diagnostic import manifest_check
    for p in (AF_MODEL,AF_EXP,ROOT/'models/male-cns-lif/author-apl-lab-v1',ROOT/'experiments/phase4ag-author-apl'):
        result[str(p.relative_to(ROOT))]=manifest_check(p,digest)
    assert all(r['pass'] for r in result.values())
    return result

def check():
    p=json.loads((EXP/'protocol/preregistered.json').read_text())
    assert p['sources']==sources(),'Source changed after preregistration'
    for row in p['inputs']:assert digest(ROOT/row['path'])==row['sha256'],row['path']
    for row in json.loads((EXP/'protocol/prepared-hashes.json').read_text()):
        assert digest(EXP/row['path'])==row['sha256'],row['path']
    return p

def budget():
    size=sum(p.stat().st_size for p in EXP.rglob('*') if p.is_file())
    assert size<CFG['output_budget_bytes'],'New output resource budget exceeded'
    return size