import asyncio
import gc
import os
import pytest
import numpy as np
import threading
import warnings
from numpy.testing import extbuild, assert_warns, IS_WASM
import sys


@pytest.fixture
def get_module(tmp_path):
    """ Add a memory policy that returns a false pointer 64 bytes into the
    actual allocation, and fill the prefix with some text. Then check at each
    memory manipulation that the prefix exists, to make sure all alloc/realloc/
    free/calloc go via the functions here.
    """
    if sys.platform.startswith('cygwin'):
        pytest.skip('link fails on cygwin')
    if IS_WASM:
        pytest.skip("Can't build module inside Wasm")
    functions = [
        ("get_default_policy", "METH_NOARGS", """
             Py_INCREF(PyDataMem_DefaultHandler);
             return PyDataMem_DefaultHandler;
         """),
        ("set_secret_data_policy", "METH_NOARGS", """
             PyObject *secret_data =
                 PyCapsule_New(&secret_data_handler, "mem_handler", NULL);
             if (secret_data == NULL) {
                 return NULL;
             }
             PyObject *old = PyDataMem_SetHandler(secret_data);
             Py_DECREF(secret_data);
             return old;
         """),
        ("set_old_policy", "METH_O", """
             PyObject *old;
             if (args != NULL && PyCapsule_CheckExact(args)) {
                 old = PyDataMem_SetHandler(args);
             }
             else {
                 old = PyDataMem_SetHandler(NULL);
             }
             return old;
         """),
        ("get_array", "METH_NOARGS", """
            char *buf = (char *)malloc(20);
            npy_intp dims[1];
            dims[0] = 20;
            PyArray_Descr *descr =  PyArray_DescrNewFromType(NPY_UINT8);
            return PyArray_NewFromDescr(&PyArray_Type, descr, 1, dims, NULL,
                                        buf, NPY_ARRAY_WRITEABLE, NULL);
         """),
        ("set_own", "METH_O", """
            if (!PyArray_Check(args)) {
                PyErr_SetString(PyExc_ValueError,
                             "need an ndarray");
                return NULL;
            }
            PyArray_ENABLEFLAGS((PyArrayObject*)args, NPY_ARRAY_OWNDATA);
            // Maybe try this too?
            // PyArray_BASE(PyArrayObject *)args) = NULL;
            Py_RETURN_NONE;
         """),
        ("get_array_with_base", "METH_NOARGS", """
            char *buf = (char *)malloc(20);
            npy_intp dims[1];
            dims[0] = 20;
            PyArray_Descr *descr =  PyArray_DescrNewFromType(NPY_UINT8);
            PyObject *arr = PyArray_NewFromDescr(&PyArray_Type, descr, 1, dims,
                                                 NULL, buf,
                                                 NPY_ARRAY_WRITEABLE, NULL);
            if (arr == NULL) return NULL;
            PyObject *obj = PyCapsule_New(buf, "buf capsule",
                                          (PyCapsule_Destructor)&warn_on_free);
            if (obj == NULL) {
                Py_DECREF(arr);
                return NULL;
            }
            if (PyArray_SetBaseObject((PyArrayObject *)arr, obj) < 0) {
                Py_DECREF(arr);
                Py_DECREF(obj);
                return NULL;
            }
            return arr;

         """),
    ]
    prologue = '''
        #define NPY_TARGET_VERSION NPY_1_22_API_VERSION
        #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
        #include <numpy/arrayobject.h>
        /*
         * This struct allows the dynamic configuration of the allocator funcs
         * of the `secret_data_allocator`. It is provided here for
         * demonstration purposes, as a valid `ctx` use-case scenario.
         */
        typedef struct {
            void *(*malloc)(size_t);
            void *(*calloc)(size_t, size_t);
            void *(*realloc)(void *, size_t);
            void (*free)(void *);
        } SecretDataAllocatorFuncs;

        NPY_NO_EXPORT void *
        shift_alloc(void *ctx, size_t sz) {
            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx;
            char *real = (char *)funcs->malloc(sz + 64);
            if (real == NULL) {
                return NULL;
            }
            snprintf(real, 64, "originally allocated %ld", (unsigned long)sz);
            return (void *)(real + 64);
        }
        NPY_NO_EXPORT void *
        shift_zero(void *ctx, size_t sz, size_t cnt) {
            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx;
            char *real = (char *)funcs->calloc(sz + 64, cnt);
            if (real == NULL) {
                return NULL;
            }
            snprintf(real, 64, "originally allocated %ld via zero",
                     (unsigned long)sz);
            return (void *)(real + 64);
        }
        NPY_NO_EXPORT void
        shift_free(void *ctx, void * p, npy_uintp sz) {
            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx;
            if (p == NULL) {
                return ;
            }
            char *real = (char *)p - 64;
            if (strncmp(real, "originally allocated", 20) != 0) {
                fprintf(stdout, "uh-oh, unmatched shift_free, "
                        "no appropriate prefix\\n");
                /* Make C runtime crash by calling free on the wrong address */
                funcs->free((char *)p + 10);
                /* funcs->free(real); */
            }
            else {
                npy_uintp i = (npy_uintp)atoi(real +20);
                if (i != sz) {
                    fprintf(stderr, "uh-oh, unmatched shift_free"
                            "(ptr, %ld) but allocated %ld\\n", sz, i);
                    /* This happens in some places, only print */
                    funcs->free(real);
                }
                else {
                    funcs->free(real);
                }
            }
        }
        NPY_NO_EXPORT void *
        shift_realloc(void *ctx, void * p, npy_uintp sz) {
            SecretDataAllocatorFuncs *funcs = (SecretDataAllocatorFuncs *)ctx;
            if (p != NULL) {
                char *real = (char *)p - 64;
                if (strncmp(real, "originally allocated", 20) != 0) {
                    fprintf(stdout, "uh-oh, unmatched shift_realloc\\n");
                    return realloc(p, sz);
                }
                return (void *)((char *)funcs->realloc(real, sz + 64) + 64);
            }
            else {
                char *real = (char *)funcs->realloc(p, sz + 64);
                if (real == NULL) {
                    return NULL;
                }
                snprintf(real, 64, "originally allocated "
                         "%ld  via realloc", (unsigned long)sz);
                return (void *)(real + 64);
            }
        }
        /* As an example, we use the standard {m|c|re}alloc/free funcs. */
        static SecretDataAllocatorFuncs secret_data_handler_ctx = {
            malloc,
            calloc,
            realloc,
            free
        };
        static PyDataMem_Handler secret_data_handler = {
            "secret_data_allocator",
            1,
            {
                &secret_data_handler_ctx, /* ctx */
                shift_alloc,              /* malloc */
                shift_zero,               /* calloc */
                shift_realloc,            /* realloc */
                shift_free                /* free */
            }
        };
        void warn_on_free(void *capsule) {
            PyErr_WarnEx(PyExc_UserWarning, "in warn_on_free", 1);
            void * obj = PyCapsule_GetPointer(capsule,
                                              PyCapsule_GetName(capsule));
            free(obj);
        };
        '''
    more_init = "import_array();"
    try:
        import mem_policy
        return mem_policy
    except ImportError:
        pass
    # if it does not exist, build and load it
    return extbuild.build_and_import_extension('mem_policy',
                                               functions,
                                               prologue=prologue,
                                               include_dirs=[np.get_include()],
                                               build_dir=tmp_path,
                                               more_init=more_init)


def test_set_policy(get_module):

    get_handler_name = np.core.multiarray.get_handler_name
    get_handler_version = np.core.multiarray.get_handler_version
    orig_policy_name = get_handler_name()

    a = np.arange(10).reshape((2, 5))  # a doesn't own its own data
    assert get_handler_name(a) is None
    assert get_handler_version(a) is None
    assert get_handler_name(a.base) == orig_policy_name
    assert get_handler_version(a.base) == 1

    orig_policy = get_module.set_secret_data_policy()

    b = np.arange(10).reshape((2, 5))  # b doesn't own its own data
    assert get_handler_name(b) is None
    assert get_handler_version(b) is None
    assert get_handler_name(b.base) == 'secret_data_allocator'
    assert get_handler_version(b.base) == 1

    if orig_policy_name == 'default_allocator':
        get_module.set_old_policy(None)  # tests PyDataMem_SetHandler(NULL)
        assert get_handler_name() == 'default_allocator'
    else:
        get_module.set_old_policy(orig_policy)
        assert get_handler_name() == orig_policy_name


def test_default_policy_singleton(get_module):
    get_handler_name = np.core.multiarray.get_handler_name

    # set the policy to default
    orig_policy = get_module.set_old_policy(None)

    assert get_handler_name() == 'default_allocator'

    # re-set the policy to default
    def_policy_1 = get_module.set_old_policy(None)

    assert get_handler_name() == 'default_allocator'

    # set the policy to original
    def_policy_2 = get_module.set_old_policy(orig_policy)

    # since default policy is a singleton,
    # these should be the same object
    assert def_policy_1 is def_policy_2 is get_module.get_default_policy()


def test_policy_propagation(get_module):
    # The memory policy goes hand-in-hand with flags.owndata

    class MyArr(np.ndarray):
        pass

    get_handler_name = np.core.multiarray.get_handler_name
    orig_policy_name = get_handler_name()
    a = np.arange(10).view(MyArr).reshape((2, 5))
    assert get_handler_name(a) is None
    assert a.flags.owndata is False

    assert get_handler_name(a.base) is None
    assert a.base.flags.owndata is False

    assert get_handler_name(a.base.base) == orig_policy_name
    assert a.base.base.flags.owndata is True


async def concurrent_context1(get_module, orig_policy_name, event):
    if orig_policy_name == 'default_allocator':
        get_module.set_secret_data_policy()
        assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
    else:
        get_module.set_old_policy(None)
        assert np.core.multiarray.get_handler_name() == 'default_allocator'
    event.set()


async def concurrent_context2(get_module, orig_policy_name, event):
    await event.wait()
    # the policy is not affected by changes in parallel contexts
    assert np.core.multiarray.get_handler_name() == orig_policy_name
    # change policy in the child context
    if orig_policy_name == 'default_allocator':
        get_module.set_secret_data_policy()
        assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
    else:
        get_module.set_old_policy(None)
        assert np.core.multiarray.get_handler_name() == 'default_allocator'


async def async_test_context_locality(get_module):
    orig_policy_name = np.core.multiarray.get_handler_name()

    event = asyncio.Event()
    # the child contexts inherit the parent policy
    concurrent_task1 = asyncio.create_task(
        concurrent_context1(get_module, orig_policy_name, event))
    concurrent_task2 = asyncio.create_task(
        concurrent_context2(get_module, orig_policy_name, event))
    await concurrent_task1
    await concurrent_task2

    # the parent context is not affected by child policy changes
    assert np.core.multiarray.get_handler_name() == orig_policy_name


def test_context_locality(get_module):
    if (sys.implementation.name == 'pypy'
            and sys.pypy_version_info[:3] < (7, 3, 6)):
        pytest.skip('no context-locality support in PyPy < 7.3.6')
    asyncio.run(async_test_context_locality(get_module))


def concurrent_thread1(get_module, event):
    get_module.set_secret_data_policy()
    assert np.core.multiarray.get_handler_name() == 'secret_data_allocator'
    event.set()


def concurrent_thread2(get_module, event):
    event.wait()
    # the policy is not affected by changes in parallel threads
    assert np.core.multiarray.get_handler_name() == 'default_allocator'
    # change policy in the child thread
    get_module.set_secret_data_policy()


def test_thread_locality(get_module):
    orig_policy_name = np.core.multiarray.get_handler_name()

    event = threading.Event()
    # the child threads do not inherit the parent policy
    concurrent_task1 = threading.Thread(target=concurrent_thread1,
                                        args=(get_module, event))
    concurrent_task2 = threading.Thread(target=concurrent_thread2,
                                        args=(get_module, event))
    concurrent_task1.start()
    concurrent_task2.start()
    concurrent_task1.join()
    concurrent_task2.join()

    # the parent thread is not affected by child policy changes
    assert np.core.multiarray.get_handler_name() == orig_policy_name


@pytest.mark.slow
def test_new_policy(get_module):
    a = np.arange(10)
    orig_policy_name = np.core.multiarray.get_handler_name(a)

    orig_policy = get_module.set_secret_data_policy()

    b = np.arange(10)
    assert np.core.multiarray.get_handler_name(b) == 'secret_data_allocator'

    # test array manipulation. This is slow
    if orig_policy_name == 'default_allocator':
        # when the np.core.test tests recurse into this test, the
        # policy will be set so this "if" will be false, preventing
        # infinite recursion
        #
        # if needed, debug this by
        # - running tests with -- -s (to not capture stdout/stderr
        # - setting extra_argv=['-vv'] here
        assert np.core.test('full', verbose=2, extra_argv=['-vv'])
        # also try the ma tests, the pickling test is quite tricky
        assert np.ma.test('full', verbose=2, extra_argv=['-vv'])

    get_module.set_old_policy(orig_policy)

    c = np.arange(10)
    assert np.core.multiarray.get_handler_name(c) == orig_policy_name

@pytest.mark.xfail(sys.implementation.name == "pypy",
                   reason=("bad interaction between getenv and "
                           "os.environ inside pytest"))
@pytest.mark.parametrize("policy", ["0", "1", None])
def test_switch_owner(get_module, policy):
    a = get_module.get_array()
    assert np.core.multiarray.get_handler_name(a) is None
    get_module.set_own(a)
    oldval = os.environ.get('NUMPY_WARN_IF_NO_MEM_POLICY', None)
    if policy is None:
        if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
            os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
    else:
        os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = policy
    try:
        # The policy should be NULL, so we have to assume we can call
        # "free".  A warning is given if the policy == "1"
        if policy == "1":
            with assert_warns(RuntimeWarning) as w:
                del a
                gc.collect()
        else:
            del a
            gc.collect()

    finally:
        if oldval is None:
            if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
                os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
        else:
            os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = oldval

def test_owner_is_base(get_module):
    a = get_module.get_array_with_base()
    with pytest.warns(UserWarning, match='warn_on_free'):
        del a
        gc.collect()