"""Tiny host-import/lifetime receipts from the pinned MTPLX runtime; no models."""
import gc
import json
import mmap
import os
import weakref

import mlx.core as mx
import numpy as np


def main():
    assert mx.__version__ == "0.32.2"
    page = os.sysconf("SC_PAGE_SIZE")
    mx.set_cache_limit(4 * page)
    for name, offset, size in (
        ("aligned", 0, page),
        ("unaligned", 1, page),
        ("short", 0, 17),
        ("empty", 0, 0),
    ):
        mx.clear_cache()
        baseline = mx.get_active_memory()
        owner = mmap.mmap(-1, 2 * page)
        source = np.ndarray((size,), dtype=np.uint8, buffer=owner, offset=offset)
        source[:] = 37
        pointer = source.__array_interface__["data"][0]
        weak = weakref.ref(source)
        # The codec calls this constructor: array.cpp passes copy=true to
        # create_array, even though nd_array_to_mlx has a different default.
        value = mx.array(source)
        mx.eval(value)
        exported = np.asarray(value)
        adopted = bool(size and pointer == exported.__array_interface__["data"][0])
        assert (exported == 37).all()
        del exported, source
        gc.collect()
        retained = weak() is not None
        assert not adopted and not retained
        active = mx.get_active_memory() - baseline
        view = value[:]
        mx.eval(view)
        del value
        gc.collect()
        retained_by_view = weak() is not None
        assert not retained_by_view
        assert (np.asarray(view) == 37).all()
        del view
        gc.collect()
        mx.synchronize()
        released = weak() is None
        assert released
        assert mx.get_active_memory() == baseline
        cached = mx.get_cache_memory()
        owner.close()
        print(json.dumps(dict(event="host_import", name=name, page=page, size=size,
                              adopted=adopted, active=active, retained=retained,
                              retained_by_view=retained_by_view, released=released,
                              cached_after_release=cached)), flush=True)


if __name__ == "__main__":
    main()