"""One pre-loop overlay hook in the SHA-pinned AJ function; no old global patch.

run_gpu_overlay(case_dir, out_dir, build_root, overlay_file, overlay_manifest)
retains AJ's output/summary schema. overlay-audit.json and adapter-provenance.json
are mandatory additional evidence. The inherited source-audit.json describes
the unchanged AJ source only, NOT the complete Phase 4AM adapter semantics.
Delta is fixed during the whole trial; this interface never trains or updates
eligibility. The original CUDA arrival operator still samples current weights
at arrival, and no emission-time weight snapshot is introduced.
"""
from __future__ import annotations

import ast
import copy
import hashlib
import importlib.util
import json
from pathlib import Path
import sys

import numpy as np

import overlay


AJ_SHA256 = 'af1975ba5e5e700208c0a598d3cbe648a3cf10885c6231239b6e4b28f76d173b'
PINS = {
    'apl-timing-lab-v1/gpu_timing.py': AJ_SHA256,
    'apl-timing-lab-v1/timing_queue.py': '54df6d4818e2a0c26e96fb0aaf5d9b71cd964981706da937562d48af79aca14e',
    'author-apl-lab-v2/gpu_apl.py': '69f5d893f47dafff4ef7c3186c3df0b9014666b32c8f4e943b57016bec92b18a',
    'author-apl-lab-v2/legacy_apl_worker.py': '21ceec65fe42581dc55872010e0025d660e80358905f1cf4df1dadf341a3b291',
    'author-runtime-lab-v1/gpu_backend.py': '03758d07572f4b44d272059d9308fac534268ad201b57bdd77d69fee95e67278',
    'author-runtime-lab-v1/author_ext.cpp': '208a937d17bbb041b9bf6b25b89264ee566830b60609e58491b44583d8541f15',
    'author-runtime-lab-v1/author_ext.cu': '669f166e18ce9a3f63fc723db8b0a8b74f7fa4757194b1bd65b4cf6008a7787e',
    'author-fullgraph-lab-v1/oracle_full.py': 'd72003d98f0f1624b12126eff50bd361c43d888f4a64d90ba9247a58976f1961',
}
HOOK_NAME = '_phase4am_install_overlay'
HOOK_CODE = HOOK_NAME + '(weights, graph, apl_out_mask, upload, out_dir, source_provenance)'
_LOADED = {}


def _save(path, value):
    with Path(path).open('x', encoding='utf-8') as stream:
        json.dump(value, stream, ensure_ascii=False, sort_keys=True, indent=2, allow_nan=False)
        stream.write('\n')


def verify_sources(aj_model_dir):
    directory = Path(aj_model_dir).absolute()
    overlay.require(directory.name == 'apl-timing-lab-v1', 'AJ must be loaded from its actual original model directory')
    paths = {key: directory.parent / key for key in PINS}
    for key, path in paths.items():
        overlay.require(overlay.file_hash(path) == PINS[key], 'Sealed source SHA mismatch: ' + str(path))
    return {key: {'path': str(paths[key]), 'sha256': value} for key, value in PINS.items()}


def _load_checked(name, path, expected):
    path = Path(path).absolute()
    data = path.read_bytes()
    overlay.require(hashlib.sha256(data).hexdigest() == expected, 'Source changed before loading: ' + str(path))
    if name in sys.modules:
        existing = sys.modules[name]
        overlay.require(_LOADED.get(name) == (path, expected, existing), 'Refusing an existing unverified module binding: ' + name)
        return existing
    spec = importlib.util.spec_from_file_location(name, path)
    overlay.require(spec is not None and spec.loader is not None, 'No source loader')
    module = importlib.util.module_from_spec(spec)
    sys.modules[name] = module
    try:
        # Execute precisely the checked source bytes; do not write frozen pyc.
        exec(compile(data, str(path), 'exec'), module.__dict__)
    except BaseException:
        sys.modules.pop(name, None)
        raise
    _LOADED[name] = (path, expected, module)
    return module


def _dump(node):
    return ast.dump(node, annotate_fields=True, include_attributes=False)


def derive_ast(source_bytes):
    """Pure AST adapter: one inserted Expr, every other original node unchanged."""
    overlay.require(hashlib.sha256(source_bytes).hexdigest() == AJ_SHA256, 'Cannot derive from unpinned AJ source')
    module = ast.parse(source_bytes)
    matches = [node for node in module.body if isinstance(node, ast.FunctionDef) and node.name == 'run_gpu']
    overlay.require(len(matches) == 1, 'Expected exactly one sealed AJ run_gpu')
    original = matches[0]
    adapted = copy.deepcopy(original)
    tries = [node for node in adapted.body if isinstance(node, ast.Try)]
    overlay.require(len(tries) == 1, 'Unexpected run_gpu try structure')
    statements = tries[0].body
    anchor = ast.parse('effective_hash = hashlib.sha256()').body[0]
    positions = [i for i, node in enumerate(statements) if _dump(node) == _dump(anchor)]
    overlay.require(len(positions) == 1, 'Effective-weight hash insertion anchor is not unique')
    position = positions[0]
    previous = ast.parse('if bool(np.any(apl_out_mask)):\n    weights[upload(apl_out_mask, torch.bool)] = 0.0').body[0]
    overlay.require(position > 0 and _dump(statements[position - 1]) == _dump(previous),
                    'Overlay hook must immediately follow frozen fast-APL zeroing')
    hook = ast.copy_location(ast.parse(HOOK_CODE).body[0], statements[position])
    statements.insert(position, hook)
    tick_loops = lambda tree: [node for node in ast.walk(tree) if isinstance(node, ast.For) and
                              isinstance(node.target, ast.Name) and node.target.id == 'tick']
    old_loops, new_loops = tick_loops(original), tick_loops(adapted)
    overlay.require(len(old_loops) == len(new_loops) == 1 and _dump(old_loops[0]) == _dump(new_loops[0]),
                    'Entire simulation loop AST must remain identical')
    restored = copy.deepcopy(adapted)
    restored_try = next(node for node in restored.body if isinstance(node, ast.Try))
    restored_try.body.pop(position)
    overlay.require(_dump(restored) == _dump(original), 'Adapter changed more than its single pre-loop hook')
    ast.fix_missing_locations(adapted)
    sha = lambda value: hashlib.sha256(value.encode('utf-8')).hexdigest()
    audit = {'original_run_ast_sha256': sha(_dump(original)), 'adapted_run_ast_sha256': sha(_dump(adapted)),
             'loop_ast_sha256': sha(_dump(old_loops[0])), 'entire_tick_loop_ast_identical': True,
             'only_ast_change': HOOK_CODE, 'inserted_ast_statements': 1,
             'hook_after_apl_zero_before_effective_hash': True}
    return adapted, audit


def compile_derived(aj_module, install_hook):
    """Compile only the derived function into a NEW namespace, never old globals."""
    source_path = Path(aj_module.__file__).absolute()
    function_ast, audit = derive_ast(source_path.read_bytes())
    namespace = dict(aj_module.__dict__)
    # Real adapter location, not AJ's __file__ masquerading as a new module.
    namespace.update(__name__=__name__ + '.derived_run', __file__=str(Path(__file__).absolute()))
    overlay.require(HOOK_NAME not in namespace, 'Hook name collides with sealed globals')
    namespace[HOOK_NAME] = install_hook
    tree = ast.fix_missing_locations(ast.Module(body=[function_ast], type_ignores=[]))
    exec(compile(tree, str(Path(__file__).absolute()) + '::AST-derived-AJ-run_gpu', 'exec'), namespace)
    audit.update(derived_namespace_file=namespace['__file__'], original_module_file=str(source_path),
                 original_module_globals_patched=False, fake_original_file_used=False)
    return namespace['run_gpu'], audit


def _copy_input(source, destination, expected):
    with Path(source).open('rb') as incoming, Path(destination).open('xb') as outgoing:
        for block in iter(lambda: incoming.read(8 << 20), b''):
            outgoing.write(block)
    overlay.require(overlay.file_hash(destination) == expected, 'Copied input hash mismatch')


def run_gpu_overlay(case_dir, out_dir, build_root, overlay_file, overlay_manifest):
    """Fresh D18 constant-weight interface run; no learning/update inside trial.

    This is the ONLY GPU entry point; importing/testing the module is CPU-only.
    Output is exclusive. The caller still owns preregistration, case/graph and
    frozen-mask authority, resource checks, budgets, and scientific gate choice.
    Success requires the added complete adapter-provenance.json, not merely an
    inherited AJ summary that may precede a failed final wrapper verification.
    """
    case_dir, out_dir = Path(case_dir).absolute(), Path(out_dir).absolute()
    overlay.require(not out_dir.exists(), 'Refusing an existing output directory')
    overlay.require(sys.dont_write_bytecode, 'Invoke with python -B so inherited loaders cannot write frozen __pycache__')
    own_sources = {'adapter': overlay.file_hash(__file__), 'overlay': overlay.file_hash(overlay.__file__)}
    plan = overlay.load_overlay(overlay_file, overlay_manifest)
    spec = json.loads((case_dir / 'spec.json').read_text(encoding='utf-8'))
    overlay.require(type(spec.get('apl_delay_steps')) is int and spec['apl_delay_steps'] == 18,
                    'Phase 4AM interface requires fixed D18')
    overlay.require(spec.get('training', False) is False and spec.get('learning', False) is False,
                    'This interface does not authorize a learning run')
    aj_dir = Path(plan['metadata']['aj_model_dir'])
    sources = verify_sources(aj_dir)
    _load_checked('timing_queue', aj_dir / 'timing_queue.py', PINS['apl-timing-lab-v1/timing_queue.py'])
    aj = _load_checked('phase4am_sealed_aj_gpu', aj_dir / 'gpu_timing.py', AJ_SHA256)
    base_file = Path(spec['graph_dir']) / 'weights.npy'
    base_file_sha = overlay.file_hash(base_file)
    calls = []

    def install(weights, graph, apl_out_mask, upload, output, inherited_provenance):
        overlay.require(not calls, 'Plastic overlay hook must run exactly once')
        mapping = overlay.map_installation(graph['weights'], graph['edge_id'], apl_out_mask, plan)
        if len(graph['nodes']) == 166700:
            overlay.require(int(apl_out_mask.sum()) == 5845, 'Frozen fullgraph fast-APL mask changed')
        torch = aj.torch
        overlay.require(weights.dtype == torch.float64 and weights.ndim == 1 and weights.numel() == len(graph['edge_id']),
                        'Runtime weight tensor must retain float64 SI edge order')
        rows = mapping['rows']
        if len(rows):
            local = upload(rows, torch.int64)
            # Separate eager f64 multiply and addition, never base->f32 roundtrip.
            delta32 = plan['plastic_delta_mv'][plan['plastic_delta_mv'] != np.float32(0)]
            delta_si = upload(delta32, torch.float64) * 0.001
            weights[local] = weights[local] + delta_si
            actual = weights[local].cpu().numpy().copy()
            overlay.require(overlay.bit_equal(actual, mapping['expected_si']), 'GPU installation differs from independent CPU SI arithmetic')
        # With zero delta no indexing assignment or +0 arithmetic happens.
        base_raw_sha, expected_sha = overlay.expected_hashes(graph['weights'], apl_out_mask, mapping)
        actual_hash = hashlib.sha256()
        for start in range(0, weights.numel(), 1 << 20):
            actual_hash.update(weights[start:start + (1 << 20)].cpu().numpy().astype('<f8', copy=False).tobytes())
        effective_sha = actual_hash.hexdigest()
        overlay.require(effective_sha == expected_sha, 'Full effective weight bytes differ from base+APLzero+sparse overlay')
        overlay.verify_inputs(plan)
        overlay.require(overlay.file_hash(base_file) == base_file_sha, 'Original graph base file changed')
        _copy_input(plan['input']['overlay_file'], output / 'overlay-input.npz', plan['input']['overlay_sha256'])
        _copy_input(plan['input']['manifest_file'], output / 'overlay-manifest.json', plan['input']['manifest_sha256'])
        audit = {'schema': 'phase4am-overlay-audit-v1', 'complete': True,
            'overlay_name': plan['metadata']['overlay_name'], 'input': plan['input'],
            'arrays': {key: overlay.array_metadata(plan[key]) for key in
                       ('edge_id', 'plastic_delta_mv', 'base_candidate_mv') if plan[key] is not None},
            'candidate_count': len(plan['edge_id']), 'nonzero_delta_count': int(np.count_nonzero(plan['plastic_delta_mv'])),
            **{key: mapping[key] for key in ('mapped_candidate_count', 'missing_candidate_count', 'installed_nonzero_count', 'apl_overlap_count')},
            'delta_dtype': 'float32', 'delta_units': 'mV', 'state_dtype': 'float64', 'bridge': overlay.BRIDGE,
            'installation_schedule': 'before tick 0; after fast APL zeroing',
            'weights_updated_during_run': False, 'learning': False, 'zero_delta_no_write': True,
            'frozen_mask_exact': True, 'bounds_pass': True, 'apl_fast_edges_zeroed': int(apl_out_mask.sum()),
            'base_graph_files_modified': False, 'untouched_weights_bitwise_preserved': True,
            'base_weights_file_sha256': base_file_sha, 'base_weights_raw_sha256': base_raw_sha,
            'effective_weights_raw_sha256': effective_sha, 'expected_effective_weights_raw_sha256': expected_sha,
            'copied_overlay_file': 'overlay-input.npz', 'copied_manifest_file': 'overlay-manifest.json',
            'inherited_source_audit_scope': 'AJ source provenance only; AM overlay semantics are recorded separately',
            'inherited_aj_source_provenance': inherited_provenance,
            'static25_is_exact_fraction_of_float64_base_claimed': False,
            'arrival_weight_sampling': 'arrival time; delta fixed for entire trial'}
        _save(output / 'overlay-audit.json', audit)
        calls.append(audit)

    derived, ast_audit = compile_derived(aj, install)
    try:
        summary = derived(case_dir, out_dir, Path(build_root))
        overlay.require(len(calls) == 1, 'Exactly one overlay installation was required')
        overlay.verify_inputs(plan)
        overlay.require(verify_sources(aj_dir) == sources and overlay.file_hash(base_file) == base_file_sha,
                        'Frozen sources or base graph changed during the run')
        overlay.require(own_sources == {'adapter': overlay.file_hash(__file__), 'overlay': overlay.file_hash(overlay.__file__)},
                        'Adapter source changed during invocation')
        overlay.require(summary.get('complete') is True and summary.get('effective_weights_raw_sha256') ==
                        calls[0]['effective_weights_raw_sha256'], 'Inherited run summary does not bind installed weights')
        receipt = {'schema': 'phase4am-adapter-provenance-v1', 'complete': True,
            'scope': 'constant sparse overlay installation only; no within-trial plasticity or learning',
            'adapter_file': str(Path(__file__).absolute()), 'adapter_sha256': own_sources['adapter'],
            'overlay_module_file': str(Path(overlay.__file__).absolute()), 'overlay_module_sha256': own_sources['overlay'],
            'sealed_sources': sources, 'ast': ast_audit, 'installation_calls': len(calls),
            'weights_updated_during_run': False, 'learning': False, 'apl_delay_steps': 18,
            'original_module_globals_patched': False, 'delta_constant_for_entire_trial': True,
            'inherited_source_audit_describes_AM_semantics': False,
            'overlay_audit_sha256': overlay.file_hash(out_dir / 'overlay-audit.json'),
            'inherited_summary_sha256': overlay.file_hash(out_dir / 'summary.json'),
            'effective_weights_raw_sha256': calls[0]['effective_weights_raw_sha256']}
        _save(out_dir / 'adapter-provenance.json', receipt)
        return summary
    except Exception as error:
        if out_dir.is_dir() and not (out_dir / 'adapter-error.json').exists():
            _save(out_dir / 'adapter-error.json', {'complete': False, 'type': type(error).__name__,
                'message': str(error), 'installation_calls': len(calls), 'no_complete_adapter_receipt': True})
        raise