#!/usr/bin/env python3
"""Preserve names without changing the original target-feature build's code.

prepare ORIGINAL_RAW NAMED_RAW OUTPUT_RAW
inspect ORIGINAL_BG NAMED_BG ORIGINAL_WAT NAMED_WAT
"""
import hashlib
import importlib.util
import json
from pathlib import Path
import re
import sys

# inspect.py beside this file owns the SIMD opcode pattern. Load it by path under
# its own module name so it never shadows the standard library's `inspect`.
_spec = importlib.util.spec_from_file_location('wat_inspect', Path(__file__).with_name('inspect.py'))
wat_inspect = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(wat_inspect)


def require(condition, message):
    """Reject and exit nonzero. Unlike `assert`, this survives `python3 -O`."""
    if not condition:
        raise SystemExit(f'attribute.py: {message}')


def leb(data, offset):
    result = shift = 0
    while True:
        value = data[offset]
        offset += 1
        result |= (value & 127) << shift
        shift += 7
        if value < 128:
            return result, offset


def sections(data):
    require(data[:8] == b'\0asm\x01\0\0\0', 'Not a Wasm binary')
    offset = 8
    while offset < len(data):
        start = offset
        kind = data[offset]
        size, begin = leb(data, offset + 1)
        offset = begin + size
        yield kind, data[begin:offset], data[start:offset]


def standard(data):
    return {kind: body for kind, body, _ in sections(data) if kind}


def custom_name(body):
    size, offset = leb(body, 0)
    return body[offset:offset + size]


def functions(wat):
    """(identifier, SIMD opcode counts) for every module-level function, in order."""
    # Binaryen prints each module-level function at one-space indentation.
    starts = list(re.finditer(r'^ \(func (\$[^\s()]+)', wat, re.M))
    result = []
    for i, match in enumerate(starts):
        end = starts[i + 1].start() if i + 1 < len(starts) else len(wat)
        result.append((match[1], wat_inspect.count_simd(wat[match.start():end])))
    return result


def function_exports(wat):
    """Export name to function identifier, from the module's export entries."""
    return dict(re.findall(r'^ \(export "([^"]+)" \(func (\$[^\s()]+)\)\)', wat, re.M))


if sys.argv[1] == 'prepare':
    original, named = [Path(p).read_bytes() for p in sys.argv[2:4]]
    require(standard(original) == standard(named), 'Raw executable sections differ: do not attribute')
    keep = {custom_name(body) for kind, body, _ in sections(original) if kind == 0}
    keep.add(b'name')
    result = named[:8] + b''.join(full for kind, body, full in sections(named)
                                  if kind or custom_name(body) in keep)
    Path(sys.argv[4]).write_bytes(result)
    print('Raw build: every standard section is identical; original custom sections plus name retained.')
elif sys.argv[1] == 'inspect':
    original, named = [standard(Path(p).read_bytes()) for p in sys.argv[2:4]]
    require(original[10] == named[10], 'Code sections differ: do not attribute')
    # Name retention lets wasm-bindgen export allocator names instead of generic
    # __wbindgen_export names. Check everything else; report export difference.
    require({k: v for k, v in original.items() if k != 7} == {k: v for k, v in named.items() if k != 7},
            'Sections other than exports differ: do not attribute')
    original_wat, named_wat = [Path(p).read_text() for p in sys.argv[4:6]]
    before, after = functions(original_wat), functions(named_wat)
    require(len(before) == len(after), f'Function counts differ: {len(before)} vs {len(after)}')
    rows = []
    for (anonymous, counts), (symbol, named_counts) in zip(before, after):
        require(counts == named_counts, f'SIMD counts differ between {anonymous} and {symbol}')
        rows.append({'id': anonymous, 'symbol': symbol, 'count': sum(counts.values()),
                     'opcodes': dict(sorted(counts.items()))})
    total = sum(row['count'] for row in rows)
    require(total == sum(wat_inspect.count_simd(original_wat).values()),
            'Per-function SIMD counts do not add up to the whole-module count')
    print(json.dumps({'code_section_bytes': len(original[10]),
                      'code_section_sha256': hashlib.sha256(original[10]).hexdigest(),
                      'export_sections_identical': original[7] == named[7],
                      'simd_instructions': total,
                      'function_count': len(rows),
                      'exports': function_exports(original_wat),
                      'functions': rows}, indent=2))
else:
    raise SystemExit(__doc__)