"""Register a two-arm timing intervention before new trajectory outcomes."""
from datetime import datetime, timezone
import difflib
import platform
import shutil
import subprocess
import sys
import numpy as np
from context import (ROOT, MODEL, EXP, AG, AG_MODEL, AI, AI_MODEL, AH_MODEL, META,
                     COEFF, CFG, digest, read, save, sources, integrity, budget)


def copy_exact(src, dst):
    dst.parent.mkdir(parents=True, exist_ok=True)
    with src.open('rb') as a, dst.open('xb') as b:
        shutil.copyfileobj(a, b)
    assert digest(src) == digest(dst)


def main():
    assert not EXP.exists() or not any(EXP.iterdir()), 'Never reuse a populated experiment'
    assert shutil.disk_usage(EXP.parent).free > CFG['output_budget_bytes'] + 10 * 1024**3
    assert read(AI / 'audit/final-status.json')['status'] == 'PASS-CONDITIONAL-PN-APL-SOURCE-AUDIT'
    assert CFG['apl'] == read(AG_MODEL / 'config.json')['apl'] and CFG['delays_steps'] == [0, 18]
    for sub in ('protocol', 'cases', 'gpu', 'oracle', 'audit', 'logs', 'regression', 'component-audit', 'calibration', 'extension-build'):
        (EXP / sub).mkdir(parents=True, exist_ok=False)
    command = [sys.executable, '-B', '-m', 'unittest', 'discover', '-v', '-s', str(MODEL), '-p', 'test_*.py']
    tested = subprocess.run(command, cwd=MODEL, capture_output=True, text=True)
    with (EXP / 'regression/unit-tests.log').open('x') as f:
        f.write(tested.stdout + tested.stderr)
    save(EXP / 'regression/unit-tests.json', dict(command=command, exit_code=tested.returncode,
         pass_all=tested.returncode == 0, sha256=digest(EXP / 'regression/unit-tests.log')))
    assert tested.returncode == 0, 'Unit tests failed; stop before registration'
    inputs = [AI / 'manifest.json', AI / 'checksums.sha256', AI_MODEL / 'manifest.json',
              AG / 'manifest.json', AG_MODEL / 'manifest.json', META, COEFF,
              AH_MODEL / 'replay_core.py', AG_MODEL / 'legacy_apl_worker.py', AG_MODEL / 'gpu_apl.py',
              AG_MODEL / 'oracle_apl.py', AG_MODEL / 'compare.py', AG_MODEL / 'calibration.py',
              AG_MODEL / 'validate_components.py', AI / 'protocol/neuron-source-groups.parquet']
    save(EXP / 'protocol/preregistered.json', dict(timestamp_utc=datetime.now(timezone.utc).isoformat(),
        sources=sources(), config=CFG, inputs=[dict(path=str(p.relative_to(ROOT)), sha256=digest(p)) for p in inputs],
        parent_commit=subprocess.check_output(['git', '-C', str(ROOT), 'rev-parse', 'HEAD'], text=True).strip(),
        before_new_trajectories=True, prior_CAL_and_AI_outcomes_known=True,
        single_factor='generated graded-APL h applied at t versus t+18; nothing else',
        fixed_input_counts_and_amplitudes=True, actual_accepted_requests_may_change_due_to_closed_loop=True,
        formal_learning_claim=False, oracle_limit='Independent Brian LIF state with shared legacy Torch APL arithmetic; separate component audit',
        delay_biological_measurement=False, whole_brain_equal_total_inhibition_claim=False,
        timing_test_no_full_future_tail_claim=True, source_diffs_are_review_material_not_equivalence_proof=True,
        scope='Origin audit -> exact shift conservation -> two-arm fixtures -> fullgraph oracle probes -> complete two-arm fixed CAL -> frozen regression',
        python=sys.version, numpy=np.__version__, machine=platform.machine()))
    save(EXP / 'regression/old-integrity-before.json', integrity())
    for old, new in (('gpu_apl.py', 'gpu_timing.py'), ('oracle_apl.py', 'oracle_timing.py'), ('compare.py', 'compare_timing.py')):
        difference = ''.join(difflib.unified_diff((AG_MODEL / old).read_text().splitlines(True),
            (MODEL / new).read_text().splitlines(True), fromfile=str(AG_MODEL / old), tofile=str(MODEL / new)))
        with (EXP / 'protocol' / (new + '.diff')).open('x') as f:
            f.write(difference)
    records = read(AG / 'protocol/calibration-cases.json')
    assert len(records) == 60 and len({r['name'] for r in records}) == 60
    copy_exact(AG / 'protocol/calibration-cases.json', EXP / 'protocol/calibration-cases.json')
    copy_exact(AG / 'protocol/encoder-audit.json', EXP / 'protocol/encoder-audit.json')
    names = CFG['fixtures'] + [r['name'] for r in records]
    inventory = []
    for delay in CFG['delays_steps']:
        arm = 'delay' + str(delay)
        for sub in ('cases', 'gpu', 'oracle', 'component-audit', 'calibration'):
            (EXP / sub / arm).mkdir(exist_ok=False)
        for name in names:
            src = AG / 'cases' / name
            dst = EXP / 'cases' / arm / name
            dst.mkdir(exist_ok=False)
            spec = read(src / 'spec.json')
            old_spec = dict(spec)
            spec.update(apl_delay_steps=delay, budget_root=str(EXP), output_budget_bytes=CFG['output_budget_bytes'])
            save(dst / 'spec.json', spec)
            for key in set(old_spec) - {'budget_root', 'output_budget_bytes'}:
                assert old_spec[key] == spec[key]
            files = []
            for fname in ('requests.npz', 'arrays.npz', 'pn-drive.npz'):
                if (src / fname).exists():
                    copy_exact(src / fname, dst / fname)
                    files.append(dict(file=fname, sha256=digest(src / fname)))
            inventory.append(dict(arm=arm, delay_steps=delay, name=name, kind=spec['kind'],
                                  source_case=str(src.relative_to(ROOT)), unchanged_input_files=files))
    save(EXP / 'protocol/case-inventory.json', inventory)
    save(EXP / 'protocol/prepared-hashes.json', [dict(path=str(p.relative_to(EXP)), sha256=digest(p), bytes=p.stat().st_size)
         for sub in ('protocol', 'cases') for p in sorted((EXP / sub).rglob('*')) if p.is_file() and p.name != 'prepared-hashes.json'])
    budget()
    print('REGISTERED 2 delays, 8 fixtures + 60 CAL each; exact original inputs; no new trajectory yet', flush=True)


if __name__ == '__main__':
    main()