"""GCC atomic operations on aligned shared-memory fields.

The libatomic ABI omits the builtin compare-exchange ``weak`` argument:
https://github.com/gcc-mirror/gcc/blob/master/libatomic/cas_n.c
"""

from __future__ import annotations

import ctypes
from functools import cache
import mmap

_ACQUIRE, _RELEASE, _ACQ_REL, _SEQ_CST = 2, 3, 4, 5


@cache
def _library() -> ctypes.CDLL:
    try:
        library = ctypes.CDLL("libatomic.so.1")
    except OSError as exc:
        raise RuntimeError("Shared memory requires libatomic.so.1 (install libatomic1)") from exc
    library.__atomic_is_lock_free.argtypes = [ctypes.c_size_t, ctypes.c_void_p]
    library.__atomic_is_lock_free.restype = ctypes.c_bool
    for width, number in ((4, ctypes.c_uint32), (8, ctypes.c_uint64)):
        if not library.__atomic_is_lock_free(width, None):
            raise RuntimeError(f"Shared memory requires lock-free {width}-byte atomics")
        load = getattr(library, f"__atomic_load_{width}")
        load.argtypes, load.restype = [ctypes.c_void_p, ctypes.c_int], number
        store = getattr(library, f"__atomic_store_{width}")
        store.argtypes, store.restype = [ctypes.c_void_p, number, ctypes.c_int], None
    library.__atomic_compare_exchange_8.argtypes = [
        ctypes.c_void_p, ctypes.POINTER(ctypes.c_uint64), ctypes.c_uint64, ctypes.c_int, ctypes.c_int,
    ]
    library.__atomic_compare_exchange_8.restype = ctypes.c_bool
    library.atomic_thread_fence.argtypes, library.atomic_thread_fence.restype = [ctypes.c_int], None
    return library


def _address(buffer: mmap.mmap, offset: int, width: int) -> int:
    if offset % width or offset < 0 or offset + width > len(buffer):
        raise ValueError("Atomic field must be aligned and inside the mapping")
    return ctypes.addressof(ctypes.c_char.from_buffer(buffer, offset))


def load(buffer: mmap.mmap, offset: int, width: int) -> int:
    return getattr(_library(), f"__atomic_load_{width}")(_address(buffer, offset, width), _ACQUIRE)


def store(buffer: mmap.mmap, offset: int, value: int, width: int) -> None:
    if not 0 <= value < 1 << (8 * width):
        raise ValueError("Atomic field value exceeds its unsigned wire range")
    getattr(_library(), f"__atomic_store_{width}")(_address(buffer, offset, width), value, _RELEASE)


def advance(buffer: mmap.mmap, offset: int, desired: int) -> None:
    """Advance a shared cursor with CAS; concurrent progress can never be undone."""
    address = _address(buffer, offset, 8)
    expected = ctypes.c_uint64(load(buffer, offset, 8))
    while expected.value < desired:
        if _library().__atomic_compare_exchange_8(address, ctypes.byref(expected), desired, _ACQ_REL, _ACQUIRE):
            break


def fence() -> None:
    """Order invalidation/copy/validation on CPUs with weak memory ordering."""
    _library().atomic_thread_fence(_SEQ_CST)