//! Observe actual CPython allocation failures before AnySystem replaces a
//! Python exception with a generic panic. Student stdout is not a result channel.
use std::ffi::c_void;
use std::mem::MaybeUninit;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;

use pyo3::ffi::{self, PyMemAllocatorDomain, PyMemAllocatorEx};
use pyo3::Python;

static INSTALLED: Once = Once::new();
static FAILED: AtomicBool = AtomicBool::new(false);

fn observe(pointer: *mut c_void, nonzero: bool) -> *mut c_void {
    if nonzero && pointer.is_null() {
        FAILED.store(true, Ordering::Relaxed);
    }
    pointer
}

extern "C" fn malloc(ctx: *mut c_void, size: usize) -> *mut c_void {
    // SAFETY: ctx points to an immutable original allocator, leaked for the
    // interpreter's entire lifetime. Calls retain CPython's allocator contract.
    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
    observe((allocator.malloc.unwrap())(allocator.ctx, size), size != 0)
}
extern "C" fn calloc(ctx: *mut c_void, count: usize, size: usize) -> *mut c_void {
    // SAFETY: same lifetime and callback contract as malloc above.
    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
    observe(
        (allocator.calloc.unwrap())(allocator.ctx, count, size),
        count != 0 && size != 0,
    )
}
extern "C" fn realloc(ctx: *mut c_void, pointer: *mut c_void, size: usize) -> *mut c_void {
    // SAFETY: same lifetime and callback contract as malloc above.
    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
    observe(
        (allocator.realloc.unwrap())(allocator.ctx, pointer, size),
        size != 0,
    )
}
extern "C" fn free(ctx: *mut c_void, pointer: *mut c_void) {
    // SAFETY: same lifetime and callback contract as malloc above.
    let allocator = unsafe { &*ctx.cast::<PyMemAllocatorEx>() };
    (allocator.free.unwrap())(allocator.ctx, pointer);
}

pub fn install() {
    Python::attach(|_| {
        INSTALLED.call_once(|| {
            for domain in [
                PyMemAllocatorDomain::PYMEM_DOMAIN_RAW,
                PyMemAllocatorDomain::PYMEM_DOMAIN_MEM,
                PyMemAllocatorDomain::PYMEM_DOMAIN_OBJ,
            ] {
                // SAFETY: installation holds the GIL before any student code runs.
                // CPython requires post-initialization hooks to wrap the previous
                // allocator. Its four callbacks are supplied by the interpreter.
                unsafe {
                    let mut original = MaybeUninit::uninit();
                    ffi::PyMem_GetAllocator(domain, original.as_mut_ptr());
                    let original = Box::into_raw(Box::new(original.assume_init()));
                    let mut wrapper = PyMemAllocatorEx {
                        ctx: original.cast(),
                        malloc: Some(malloc),
                        calloc: Some(calloc),
                        realloc: Some(realloc),
                        free: Some(free),
                    };
                    ffi::PyMem_SetAllocator(domain, &mut wrapper);
                }
            }
        })
    });
}

pub fn reset() {
    FAILED.store(false, Ordering::Relaxed);
}
pub fn failed() -> bool {
    FAILED.load(Ordering::Relaxed)
}