"""Build a compact review bundle after all frozen endpoints and audits.

Full trajectories remain local; candidate regrading is explicitly weaker than
full raw re-extraction. No model weights or credentials are bundled.
"""
import argparse
import hashlib
import json
from pathlib import Path
import shutil
import zipfile

ROOT = Path(__file__).resolve().parents[1]
LIMIT = 2_000_000


def sha(path):
    return hashlib.sha256(path.read_bytes()).hexdigest()


def read(path):
    return json.loads(path.read_text())


def rows(path):
    return [json.loads(s) for s in path.read_text().splitlines()]


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--output', default='publication/stable_contextual_evidence_v1')
    args = parser.parse_args()
    contract = read(ROOT / 'notes/stable_contextual_contract.json')
    run = ROOT / contract['output']
    queue = read(ROOT / contract['queue_state'])
    if queue['status'] != 'COMPLETE':
        raise SystemExit('Full frozen audit is not complete; evidence packaging deferred.')
    result = read(run / 'analysis/result.json')
    assert result['status'] == 'COMPLETE_AUDITED'
    report = read(run / 'analysis/descriptive_report.json')
    assert report['primary'] == result
    out = ROOT / args.output
    assert not out.exists(), 'Never overwrite an existing evidence bundle.'
    out.mkdir(parents=True)
    original = {}

    def write(name, obj):
        p = out / name
        p.parent.mkdir(parents=True, exist_ok=True)
        p.write_text(json.dumps(obj, separators=(',', ':')) + '\n')

    def copy(source, target=None):
        p = ROOT / source
        dest = out / (target or source)
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.copyfile(p, dest)
        original[source] = sha(p)

    for source, expected in contract['input_sha256'].items():
        original[source] = expected
        if source in [contract['parent'], 'code/data/math_train.json', 'code/data/tokenizer-base.json']:
            continue
        assert sha(ROOT / source) == expected
        copy(source)
    for source in ['notes/stable_contextual_contract.json', 'notes/stable_contextual_reporting_plan.md',
                   'notes/stable_contextual_reporting_plan.md',
                   'notes/stable_contextual_environment_observed.json',
                   'code/vendor/LICENSE', 'code/vendor/reasoning_from_scratch/__init__.py',
                   'code/vendor/reasoning_from_scratch/utils.py',
                   'code/vendor/reasoning_from_scratch/ch02.py',
                   'code/requirements.txt', 'code/account_stable_contextual_work.py',
                   'code/report_stable_contextual.py', 'code/build_stable_contextual_evidence.py',
                   'code/check_stable_contextual_evidence.py']:
        copy(source)
    # Baseline warm ancestry records are distinct from this experiment's work.
    for name in ['config.json', 'training_summary.json']:
        copy(f'code/results/baseline_s37041_u50/{name}', f'ancestry/{name}')
    for name in ['result.json', 'descriptive_report.json', 'descriptive_report.md']:
        copy(str((run / 'analysis' / name).relative_to(ROOT)), 'analysis/' + name)
    train, native, snippets, groups = [], [], [], []
    for arm in contract['arms']:
        for chunk in range(contract['chunks'][arm]):
            folder = run / arm
            for name in [f'chunk{chunk}_summary.json']:
                copy(str((folder / name).relative_to(ROOT)), f'training/{arm}/{name}')
            rawpath = folder / f'chunk{chunk}_rollouts.jsonl'
            raw = rows(rawpath)
            original[str(rawpath.relative_to(ROOT))] = sha(rawpath)
            for r in raw:
                train.append([arm, r['step'], r['position'], r['sample'], r['index'],
                              len(r['tokens']), int(r['eos']), r['prediction'], int(r['reward'])])
            snippets.append({'kind': 'train', 'arm': arm, 'chunk': chunk, 'record': raw[0]})
            updatepath = folder / f'chunk{chunk}_updates.jsonl'
            original[str(updatepath.relative_to(ROOT))] = sha(updatepath)
            for update in rows(updatepath):
                for pos, g in enumerate(update['groups']):
                    groups.append([arm, update['step'], pos, g['donor'], g['receivers'],
                                   g['dense_tokens'], g['backward_calls'], g['advantages']])
            copy(f'logs/stable_contextual_train_{arm}_{chunk}.log')
        for phase in ['before', 'after']:
            copy(str((run / arm / f'retention_{phase}.json').relative_to(ROOT)),
                 f'training/{arm}/retention_{phase}.json')
    answers = {(r['arm'], r['index']): r for r in read(run / 'analysis/answers.json')}
    for arm in ['W'] + contract['arms']:
        for chunk in range(2):
            folder = run / ('eval_' + arm)
            rawpath = folder / f'chunk{chunk}_records.jsonl'
            original[str(rawpath.relative_to(ROOT))] = sha(rawpath)
            raw = rows(rawpath)
            for r in raw:
                a = answers[arm, r['index']]
                native.append([arm, r['index'], len(r['tokens']), a['prediction'], int(a['correct'])])
            snippets.append({'kind': 'eval', 'arm': arm, 'chunk': chunk, 'record': raw[0]})
            copy(str((folder / f'chunk{chunk}_summary.json').relative_to(ROOT)),
                 f'evaluation/{arm}/chunk{chunk}_summary.json')
            copy(f'logs/stable_contextual_eval_{arm}_{chunk}.log')
    assert len(train) == 6656 and len(native) == 2048 and len(snippets) == 21
    write('evidence/training.json', {'columns': ['arm','step','position','sample','index','length','eos','prediction','reward'], 'rows': train})
    write('evidence/groups.json', {'columns': ['arm','step','position','donor','receivers','dense_tokens','backward_calls','advantages'], 'rows': groups})
    write('evidence/native.json', {'columns': ['arm','index','length','prediction','correct'], 'rows': native})
    write('evidence/raw_samples.json', snippets)
    for path in sorted((ROOT / 'notes').glob('stable_contextual_interruption_*/interruption.json')):
        copy(str(path.relative_to(ROOT)))
        interruption = read(path)
        if interruption.get('stage') == 'eval':
            copy(interruption['allocation_receipt'])
            copy(interruption['retained_records'])
    write('ORIGINAL_INPUTS_AND_RAW_SHA256.json', original)
    copy('notes/stable_contextual_evidence_readme.md', 'README.md')
    files = sorted(p for p in out.rglob('*') if p.is_file())
    write('MANIFEST.json', {str(p.relative_to(out)): {'sha256':sha(p), 'bytes':p.stat().st_size} for p in files})
    expanded = sum(p.stat().st_size for p in out.rglob('*') if p.is_file())
    assert expanded < LIMIT, f'Expanded bundle {expanded} exceeds limit; preserve and inspect, do not silently omit outcomes.'
    dest = out.with_suffix('.zip')
    assert not dest.exists()
    with zipfile.ZipFile(dest, 'x', compression=zipfile.ZIP_DEFLATED) as z:
        for p in sorted(out.rglob('*')):
            if p.is_file(): z.write(p, p.relative_to(out))
    assert dest.stat().st_size < LIMIT
    print(json.dumps({'expanded_bytes': expanded, 'zip_bytes': dest.stat().st_size,
                      'zip_sha256': sha(dest), 'directory':str(out)}))


if __name__ == '__main__':
    main()
