#!/usr/bin/env python3
"""Guard the pymalloc setup block's C89 declarations and allocator order."""

import re
import sys
from pathlib import Path

source = Path(sys.argv[1]).read_text()
setup = source[source.index("\n_PyMem_SetupAllocators("):]
match = re.search(
    r"case PYMEM_ALLOCATOR_PYMALLOC:\s*"
    r"case PYMEM_ALLOCATOR_PYMALLOC_DEBUG:\s*\{(.*?)\n    \}",
    setup,
    re.S,
)
assert match, "pymalloc setup block not found"
body = match.group(1)
declarations = """
    PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
    PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;
"""
assert body.split()[:len(declarations.split())] == declarations.split(), (
    "both allocator initializers must precede the first statement for C89"
)
expected = declarations + """
    PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &malloc_alloc);
    PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &pymalloc);
    PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &pymalloc);
    if (allocator == PYMEM_ALLOCATOR_PYMALLOC_DEBUG) {
        PyMem_SetupDebugHooks();
    }
    break;
"""
assert body.split() == expected.split(), (
    "preserve allocator initializers, RAW/MEM/OBJ registration order, and debug hooks"
)
for name, values in (
    ("MALLOC_ALLOC", "NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree"),
    ("PYMALLOC_ALLOC", "NULL, _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free"),
):
    macro = re.search(r"#\s*define\s+" + name + r"\s+\{([^}]+)\}", source)
    assert macro and macro.group(1).split() == values.split(), (
        name + " must preserve ctx/malloc/calloc/realloc/free values and order"
    )
print("obmalloc setup allocators C89 declarations, field values, and RAW/MEM/OBJ/debug order: PASS")