"""Exclusive small-evidence delivery of a completed, sealed AN engineering gate.

The two external hashes are mandatory. This tool verifies bytes and declared
completion, not numerical science, and never imports an AN execution module.
Only the 21 registered source snapshots are delivered; the mutable model tree
is neither copied nor declared sealed. Failed/incomplete engineering results
retain their exact status. Partial publications are preserved, never resumed.
"""
from __future__ import annotations

import argparse
import hashlib
import importlib.util
import json
from pathlib import Path, PurePosixPath
import re
import sys


ROOT = Path('/data/digital-drosophila')
MODEL_REL = Path('models/male-cns-lif/olfactory-conditioning-lab-v1')
EXP_REL = Path('experiments/phase4an-engineering-validation')
REGISTRY_REL = Path('registry/phase4an-engineering-validation')
HELPER_REL = Path('models/male-cns-lif/pn-recruitment-audit-v1/publication.py')
HELPER_SHA = '7a79c499098b74163e7bb16adabcbc54e7069dafa204d0b99c7f4c03d79f8016'
PREREG_SHA = '7317441126f191ccc6e6d28438870b3921537ebe282a6acb163716a690057945'
MODULES = ('runtime_light', 'runtime_continuation', 'engineering_inputs', 'engineering_plan',
           'engineering_compare', 'engineering_queue_audit', 'engineering_prepare',
           'engineering_workflow', 'engineering_scopes', 'engineering_analysis')
SOURCE_FILES = tuple(sorted([m + '.py' for m in MODULES] + ['test_' + m + '.py' for m in MODULES] + ['ENGINEERING_GATE_DRAFT.md']))
PASS = 'PASS-EXACT-LIGHT-CONTINUATION-ENGINEERING-NOT-LEARNING'
STATUSES = frozenset((PASS, 'FAIL-NUMERICAL', 'INCOMPLETE-COVERAGE'))
SUFFIXES = frozenset(('.json', '.csv', '.md', '.log', '.sha256'))
MAX_BYTES = 4 * 1024**2
REQUIRED = ('protocol/preregistered.json', 'audit/final-status.json', 'audit/engineering-analysis.json',
            'audit/analysis-execution.json', 'audit/execution-complete.json',
            'regression/old-integrity-before.json', 'regression/old-integrity-after.json',
            'regression/AM-seals-after.json')


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


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


def is_sha(value):
    return type(value) is str and re.fullmatch('[0-9a-f]{64}', value) is not None


def load_helper(path):
    raw = Path(path).read_bytes()
    require(hashlib.sha256(raw).hexdigest() == HELPER_SHA, 'AI publication helper pin mismatch')
    spec = importlib.util.spec_from_file_location('phase4an_delivery_pinned_AI', path)
    module = importlib.util.module_from_spec(spec)
    exec(compile(raw, str(path), 'exec'), module.__dict__)
    return module


def _small_json(helper, experiment, rows, name):
    require(name in rows and rows[name]['bytes'] <= MAX_BYTES, 'Required small evidence missing/oversize: ' + name)
    path = helper._regular_file(experiment, name)
    require(digest(path) == rows[name]['sha256'], 'Evidence changed: ' + name)
    return helper._json_bytes(path.read_bytes())


def _write(path, value):
    raw = (json.dumps(value, sort_keys=True, indent=2, ensure_ascii=False, allow_nan=False) + '\n').encode('utf-8')
    require(len(raw) <= MAX_BYTES, 'Generated delivery evidence exceeds 4 MiB')
    with Path(path).open('xb') as stream:
        stream.write(raw)


def _sources(helper, root, experiment, rows, reg):
    model = root / MODEL_REL
    expected_paths = {str(model / name) for name in SOURCE_FILES}
    require(type(reg.get('source_sha256')) is dict and set(reg['source_sha256']) == expected_paths,
            'Registered source inventory is not the exact 21 used files')
    actual_snapshots = {p for p in rows if p.startswith('protocol/source-snapshot/')}
    require(actual_snapshots == {'protocol/source-snapshot/' + name for name in SOURCE_FILES}, 'Snapshot inventory differs from the 21 used sources')
    proof = []
    for name in SOURCE_FILES:
        relative = 'protocol/source-snapshot/' + name
        original = helper._regular_file(root, (MODEL_REL / name).as_posix())
        snapshot = helper._regular_file(experiment, relative)
        expected = reg['source_sha256'][str(original)]
        require(is_sha(expected) and rows[relative]['sha256'] == expected and
                reg['input_sha256'].get(str(snapshot)) == expected and digest(original) == digest(snapshot) == expected,
                'Actual used source / registered snapshot mismatch: ' + name)
        require(original.stat().st_size == rows[relative]['bytes'] <= MAX_BYTES, 'Used source exceeds small delivery limit')
        proof.append(dict(model_source=(MODEL_REL / name).as_posix(), sealed_source=(EXP_REL / relative).as_posix(),
                          copy='source-snapshot/' + name, bytes=rows[relative]['bytes'], sha256=expected))
    return proof


def _declared_evidence(helper, experiment, rows):
    evidence = {name: _small_json(helper, experiment, rows, name) for name in REQUIRED}
    reg = evidence['protocol/preregistered.json']; final = evidence['audit/final-status.json']
    analysis = evidence['audit/engineering-analysis.json']; execution = evidence['audit/execution-complete.json']
    require(reg.get('schema') == 'phase4an-engineering-registration-v1' and reg.get('training') is False, 'Wrong registration scope')
    require(final.get('schema') == 'phase4an-engineering-final-v1' and final.get('status') in STATUSES and
            final.get('registration_sha256') == PREREG_SHA and final.get('analysis_recomputed_exactly') is True,
            'Missing finalized/recomputed engineering evidence')
    require(analysis.get('schema') == 'phase4an-engineering-analysis-v1' and analysis.get('status') == final['status'] and
            analysis.get('registration_sha256') == PREREG_SHA and analysis.get('source_sha256') == reg['source_sha256'],
            'Final/analysis/registration disagreement')
    for obj in (final, analysis, execution):
        require(obj.get('training') is False and obj.get('learning_pass_claimed') is False, 'Unexpected learning/training claim')
    require(final.get('learned_candidate') is False and analysis.get('candidate_generated') is False and
            analysis.get('GPU_used') is False and analysis.get('no_new_brain_runs_during_analysis') is True,
            'Unexpected candidate or analysis scope')
    require(final.get('model_directory_sealed') is False and final.get('used_engineering_source_snapshots_sealed') is True,
            'Do not declare the mutable model directory sealed')
    for obj, key, expected in ((final, 'completed_calls', 32), (final, 'executed_ticks', 11074),
                              (analysis, 'expected_calls', 32), (analysis, 'analyzed_calls', 32), (analysis, 'executed_ticks', 11074),
                              (execution, 'expected_calls', 32), (execution, 'executed_ticks', 11074)):
        require(type(obj.get(key)) is int and obj[key] == expected, 'Incomplete declared engineering inventory: ' + key)
    require(all(type(analysis.get(k)) is bool for k in ('pass_all', 'numerical_exact', 'nonvacuous_temporal_coverage')),
            'Missing explicit numerical/coverage declarations')
    numerical, covered = analysis['numerical_exact'], analysis['nonvacuous_temporal_coverage']
    declared = 'FAIL-NUMERICAL' if not numerical else (PASS if covered else 'INCOMPLETE-COVERAGE')
    require(final['status'] == declared and analysis['pass_all'] is (numerical and covered), 'Declared status contradicts its numerical/coverage flags')
    analysis_sha = rows['audit/engineering-analysis.json']['sha256']
    analysis_exec = evidence['audit/analysis-execution.json']
    require(final.get('analysis_sha256') == analysis_sha and analysis_exec.get('analysis_sha256') == analysis_sha and
            analysis_exec.get('registration_sha256') == PREREG_SHA, 'Final analysis SHA linkage mismatch')
    plan = reg.get('plan', {}); planned = plan.get('rows')
    require(plan.get('total_calls') == 32 and plan.get('total_ticks') == 11074 and type(planned) is list and len(planned) == 32,
            'Incomplete registered plan')
    ids = [row.get('row_id') for row in planned]
    require(all(type(x) is str and re.fullmatch('[a-zA-Z0-9_-]+', x) for x in ids) and len(set(ids)) == 32 and
            all(type(row.get('steps')) is int and row['steps'] > 0 for row in planned) and sum(row['steps'] for row in planned) == 11074,
            'Invalid registered call identities/steps')
    complete = execution.get('completed')
    require(execution.get('complete') is True and execution.get('registration_sha256') == PREREG_SHA and
            type(complete) is list and len(complete) == 32, 'Incomplete execution receipt inventory')
    for row, done in zip(planned, complete):
        identity = row['row_id']; name = 'receipts/' + identity + '.json'
        receipt = _small_json(helper, experiment, rows, name)
        runtime_name = 'receipts/' + identity + '.runtime.json'
        runtime = _small_json(helper, experiment, rows, runtime_name)
        require(done.get('row_id') == identity and done.get('execution_sha256') == rows[name]['sha256'], 'Execution order/receipt SHA mismatch')
        require(receipt.get('complete') is True and type(receipt.get('process_exit_code')) is int and receipt['process_exit_code'] == 0 and
                receipt.get('row_id') == identity and receipt.get('registration_sha256') == PREREG_SHA and
                receipt.get('runtime_record_sha256') == rows[runtime_name]['sha256'], 'Unsuccessful/unbound process receipt')
        require(runtime.get('complete') is True and runtime.get('training') is False and runtime.get('registration_sha256') == PREREG_SHA and
                runtime.get('row') == row, 'Runtime row does not match its registered descriptor')
        for field, suffix in (('stdout_sha256', '.stdout.log'), ('stderr_sha256', '.stderr.log')):
            log = 'logs/' + identity + suffix
            require(log in rows and receipt.get(field) == rows[log]['sha256'], 'Execution log linkage mismatch')
        checkpoint = 'runs/' + identity + '/checkpoint-manifest.json'
        require(checkpoint in rows and receipt.get('checkpoint_manifest_sha256') == runtime.get('checkpoint_manifest_sha256') == rows[checkpoint]['sha256'],
                'Execution checkpoint linkage mismatch')
    before, after = (evidence['regression/old-integrity-' + point + '.json'] for point in ('before', 'after'))
    require(type(before) is dict and type(after) is dict and len(before) == len(after) == 28 and set(before) == set(after) and
            all(type(r) is dict and r.get('pass') is True for r in [*before.values(), *after.values()]) and
            final.get('old_integrity_pass') is True and final.get('old_integrity_roots') == 28 and final.get('AM_seals_preserved') is True,
            'Finalized frozen-integrity evidence missing/failed')
    for name in before:
        for report in (before[name], after[name]):
            require(report.get('mismatches', []) == [], 'Frozen integrity report contains mismatches')
        for key in ('manifest_sha256', 'checksums_sha256', 'reference_sha256'):
            if key in before[name] or key in after[name]:
                require(is_sha(before[name].get(key)) and before[name][key] == after[name].get(key),
                        'Frozen integrity authority changed: ' + name)
    am = evidence['regression/AM-seals-after.json']
    require(type(am) is dict and set(am) == {'model', 'experiment'} and all(am[k].get('pass_all') is True and
            am[k].get('manifest_sha256') == reg['anchors']['am_' + k + '_manifest_sha256'] for k in am), 'AM seal evidence contradicts registered anchors')
    return reg, final, analysis


def selection(row):
    """Return a canonical destination or an explicit omitted reason."""
    path = PurePosixPath(row['path'])
    if path.parts[0] == 'extension-build':
        return None, 'build_artifact'
    if row['bytes'] > MAX_BYTES:
        return None, 'exceeds_4_MiB'
    if path.parts[:2] == ('protocol', 'source-snapshot') and len(path.parts) == 3 and path.name in SOURCE_FILES:
        return 'source-snapshot/' + path.name, None
    if path.suffix.lower() not in SUFFIXES:
        return None, 'bulk_binary_or_unselected_suffix_including_jsonl'
    return 'evidence/' + row['path'], None


def publish(experiment_manifest_sha, final_sha, *, root=ROOT):
    require(sys.dont_write_bytecode, 'Use python -B; registered sources remain read-only')
    require(is_sha(experiment_manifest_sha) and is_sha(final_sha), 'Two actual external manifest/final SHA-256 anchors are mandatory')
    root = Path(root)
    require(root.is_absolute() and root.is_dir() and not root.is_symlink(), 'Invalid project root')
    root = root.resolve(strict=True)
    publisher_sha = digest(__file__)
    experiment, registry = root / EXP_REL, root / REGISTRY_REL
    for relative in (EXP_REL, MODEL_REL, REGISTRY_REL, HELPER_REL):
        current = root
        for part in relative.parts:
            current = current / part
            require(not current.is_symlink(), 'Symlink in publication path')
    require(not registry.exists(), 'Registry exists; no overwrite or retry')
    helper_path = root / HELPER_REL; helper = load_helper(helper_path)
    require(digest(experiment / 'manifest.json') == experiment_manifest_sha and
            digest(experiment / 'audit/final-status.json') == final_sha and
            digest(experiment / 'protocol/preregistered.json') == PREREG_SHA, 'External/fixed authority mismatch')
    verified = helper.verify_seal(experiment)
    require(verified['manifest_sha256'] == experiment_manifest_sha, 'Experiment authority changed during verification')
    rows = {row['path']: row for row in verified['files']}
    require(not any(PurePosixPath(name).name in ('stopped.json', 'preparation-failure.json', 'error.json', 'runtime-error.json') for name in rows),
            'Execution/error partial cannot be published as a completed engineering gate')
    reg, final, analysis = _declared_evidence(helper, experiment, rows)
    source_proof = _sources(helper, root, experiment, rows, reg)
    plans, omitted = [], []
    for row in verified['files'] + helper._control_rows(experiment, verified):
        destination, reason = selection(row)
        if destination:
            plans.append((row, destination))
        else:
            omitted.append(dict(row, reason=reason))
    targets = {destination for _, destination in plans}
    require(len(targets) == len(plans) and all('evidence/' + name in targets for name in (*REQUIRED, 'manifest.json', 'checksums.sha256')),
            'Required controls/evidence cannot fit the fixed small-copy policy')
    require(all('source-snapshot/' + name in targets for name in SOURCE_FILES), 'Missing deliverable registered source')
    registry.parent.mkdir(parents=True, exist_ok=True); registry.mkdir(exist_ok=False)
    try:
        records = []
        for row, destination in plans:
            source = helper._regular_file(experiment, row['path'])
            helper._copy_verified(source, registry / destination, row['bytes'], row['sha256'])
            records.append(dict(source=(EXP_REL / row['path']).as_posix(), copy=destination, bytes=row['bytes'], sha256=row['sha256']))
        require(helper.verify_seal(experiment) == verified and _sources(helper, root, experiment, rows, reg) == source_proof and
                digest(helper_path) == HELPER_SHA and digest(__file__) == publisher_sha,
                'Sealed payload, helper, publisher or original used sources changed during publication')
        _write(registry / 'seal-verification.json', {k: v for k, v in verified.items() if k != 'files'})
        _write(registry / 'evidence-index.json', records)
        _write(registry / 'source-identity.json', source_proof)
        files = [dict(path=path.relative_to(registry).as_posix(), bytes=path.stat().st_size, sha256=digest(path))
                 for path in sorted(registry.rglob('*')) if path.is_file()]
        receipt = dict(schema='phase4an-engineering-publication-v1', publication_complete=True,
            registry=REGISTRY_REL.as_posix(), project_root=str(root), declared_status=final['status'],
            declared_numerical_exact=analysis['numerical_exact'], declared_live_coverage=analysis['nonvacuous_temporal_coverage'],
            declared_engineering_pass=analysis['pass_all'], completed_calls=32, executed_ticks=11074,
            experiment_manifest_sha256=experiment_manifest_sha, experiment_checksums_sha256=verified['checksums_sha256'],
            final_sha256=final_sha, analysis_sha256=final['analysis_sha256'], preregistration_sha256=PREREG_SHA,
            publisher_sha256=publisher_sha, helper_sha256=HELPER_SHA, helper_globals_modified=False,
            used_sources_verified_before_and_after=True, used_source_count=21, model_directory_sealed=False,
            model_files_copied=False, copied_evidence_files=len(records), copied_evidence_bytes=sum(r['bytes'] for r in records),
            source_identity='source-identity.json', files=files, experiment_payload_not_copied=omitted,
            selection_policy=dict(max_file_bytes=MAX_BYTES, suffixes=sorted(SUFFIXES),
                source_snapshot_names=list(SOURCE_FILES), excluded_tree='extension-build', all_jsonl_excluded=True),
            complete_checksums_reference_bulk_not_copied=True, bulk_remains_at=EXP_REL.as_posix(), complete_bulk_backup=False,
            numerical_gate_rerun_by_publisher=False, publication_success_is_not_engineering_or_learning_pass=True,
            training=False, candidate_generated=False, learning_pass_claimed=False, existing_files_modified=False,
            receipt_self_hash='Excluded from files; its SHA is returned externally, never embedded in itself')
        path = registry / 'publication-receipt.json'; _write(path, receipt)
        return dict(publication_complete=True, declared_status=final['status'], registry=str(registry),
                    copied_evidence_files=len(records), receipt_sha256=digest(path))
    except Exception as error:
        failure = registry / 'publication-failure.json'
        if not failure.exists():
            _write(failure, dict(publication_complete=False, error_type=type(error).__name__, error=str(error), partial_outputs_preserved=True))
        raise


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--experiment-manifest-sha256', required=True)
    parser.add_argument('--final-sha256', required=True)
    args = parser.parse_args()
    print(json.dumps(publish(args.experiment_manifest_sha256, args.final_sha256), ensure_ascii=False), flush=True)


if __name__ == '__main__':
    main()