#!/usr/bin/env python3
"""Compile the exact wide-string helper as C89 and check operation ordering."""

import re
import subprocess
import sys
import tempfile
from pathlib import Path

source = Path(sys.argv[1]).read_text()
allocated = re.search(
    r"\n_Py_GetAllocatedBlocks\(void\)\s*\{(.*?)\n\}", source, re.S,
)
assert allocated, "_Py_GetAllocatedBlocks definition not found"
body = re.sub(r"/\*.*?\*/", "", allocated.group(1), flags=re.S)
assert body.split() == """
    Py_ssize_t n = raw_allocated_blocks;
    uint i;
    uintptr_t base;
    for (i = 0; i < maxarenas; ++i) {
        if (arenas[i].address == 0) {
            continue;
        }
        base = (uintptr_t)_Py_ALIGN_UP(arenas[i].address, POOL_SIZE);
        assert(base <= (uintptr_t) arenas[i].pool_address);
        for (; base < (uintptr_t) arenas[i].pool_address; base += POOL_SIZE) {
            poolp p = (poolp)base;
            n += p->ref.count;
        }
    }
    return n;
""".split(), (
    "_Py_GetAllocatedBlocks must declare uint i and uintptr_t base at entry and preserve loop operations"
)
print("obmalloc allocated-blocks C89 loop declaration and operation order: PASS")
assert re.search(
    r"_PyMem_Strdup\(const char \*str\)\s*\{\s*"
    r"size_t size;\s*char \*copy;\s*"
    r"assert\(str != NULL\);\s*size = strlen\(str\) \+ 1;\s*"
    r"copy = PyMem_Malloc\(size\);\s*"
    r"if \(copy == NULL\) \{\s*return NULL;\s*\}\s*"
    r"memcpy\(copy, str, size\);\s*return copy;\s*\}",
    source,
), "_PyMem_Strdup must use C89 entry declarations and preserve operation order"
assert re.search(
    r"_PyMem_RawStrdup\(const char \*str\)\s*\{\s*"
    r"size_t size;\s*char \*copy;\s*"
    r"assert\(str != NULL\);\s*size = strlen\(str\) \+ 1;\s*"
    r"copy = PyMem_RawMalloc\(size\);\s*"
    r"if \(copy == NULL\) \{\s*return NULL;\s*\}\s*"
    r"memcpy\(copy, str, size\);\s*return copy;\s*\}",
    source,
), "_PyMem_RawStrdup must use C89 declarations and preserve operation order"
match = re.search(
    r"wchar_t\*\s*\n_PyMem_RawWcsdup\s*\([^)]*\)\s*\{.*?\n\}",
    source, re.S,
)
assert match, "_PyMem_RawWcsdup definition not found"
harness = r'''
#include <assert.h>
#include <stddef.h>
#include <string.h>
#include <wchar.h>
#define PY_SSIZE_T_MAX 32767
static size_t length_value, allocated_size;
static int stage, fail_alloc;
static wchar_t storage[8];
static size_t checked_wcslen(const wchar_t *str) {
    assert(str != NULL && stage == 0);
    stage = 1;
    return length_value;
}
static void *PyMem_RawMalloc(size_t size) {
    assert(stage == 1);
    stage = 2;
    allocated_size = size;
    return fail_alloc ? NULL : storage;
}
static void *checked_memcpy(void *dst, const void *src, size_t size) {
    assert(stage == 2 && size == allocated_size);
    stage = 3;
    return memcpy(dst, src, size);
}
#define wcslen checked_wcslen
#define memcpy checked_memcpy
'''
harness += match.group(0)
harness += r'''
int main(void) {
    const wchar_t input[] = L"abc";
    wchar_t *result;
    length_value = 3;
    result = _PyMem_RawWcsdup(input);
    assert(result == storage && stage == 3);
    assert(allocated_size == 4 * sizeof(wchar_t));
    assert(memcmp(result, input, allocated_size) == 0);
    stage = 0;
    length_value = 0;
    result = _PyMem_RawWcsdup(L"");
    assert(result == storage && stage == 3 && storage[0] == 0);
    assert(allocated_size == sizeof(wchar_t));
    stage = 0;
    fail_alloc = 1;
    length_value = (size_t)PY_SSIZE_T_MAX / sizeof(wchar_t) - 1;
    assert(_PyMem_RawWcsdup(input) == NULL && stage == 2);
    assert(allocated_size == (length_value + 1) * sizeof(wchar_t));
    stage = 0;
    length_value++;
    assert(_PyMem_RawWcsdup(input) == NULL && stage == 1);
    return 0;
}
'''
with tempfile.TemporaryDirectory(prefix="vc6-rawwcsdup-") as directory:
    path = Path(directory)
    (path / "test.c").write_text(harness)
    subprocess.run([
        "cc", "-std=c89", "-pedantic-errors", "-Wall", "-Wextra", "-Werror",
        str(path / "test.c"), "-o", str(path / "test"),
    ], check=True)
    subprocess.run([str(path / "test")], check=True)
print("obmalloc RawWcsdup C89 compile and copy/empty/allocation-failure/overflow order: PASS")