"""
Unit tests for memory management functionality in nvshmem.core
"""
try:
    import torch
    from torch import float32
    _torch_enabled = True
except:
    float32 = None
    _torch_enabled = False

try:
    import cupy
    _cupy_enabled = True
except:
    _cupy_enabled = False

from utils import uid_init, mpi_init
import argparse
import os

import nvshmem.core

# HACK! Do not do this!
from nvshmem.core._internal_tracking import _mr_references

from cuda.core.experimental import Device, system

from mpi4py import MPI

def test_buffer():
    print("Testing basic buffer")

    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    buf = nvshmem.core.buffer(1024)
    nvshmem.core.free(buf)
    print("End basic buffer test")

def test_peer_buffer():
    print("Testing peer buffer")
    # TODO: TEAM_NODE not showing up because of anonymous enum
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    buf = nvshmem.core.buffer(1024)

    # PE0 calls get_peer_buffer on the Buffer
    if local_rank_per_node == 0:
        try:
            peer_buf = nvshmem.core.get_peer_buffer(buf, ((nvshmem.core.my_pe() + 1) % nvshmem.core.n_pes()))
        except:
            peer_buf = None
            # If >1 node (NVL domain) exists, this is an error condition. TODO enhance this to differentiate
            print("peer_buf failed to create. Check to make sure it failed because of >1 node")
            peer_buf = None
        print(peer_buf, peer_buf._mnff.ptr, buf, buf._mnff.ptr)
        if peer_buf:
            print(peer_buf, peer_buf._mnff.ptr)
        # Don't need to call free on a peer buf.
        # However, it is safe to do so. nvshmem.core.free knows to skip nvshmem_free()
        #nvshmem.core.free(peer_buf)

    nvshmem.core.free(buf)
    print("End peer buffer test")

def test_interop_torch():
    print("Testing Torch interop buffer")
    if not _torch_enabled:
        print("WARNING: Torch not found. Not running Torch Interop test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    # For the Torch test, we 
    tensor = nvshmem.core.tensor((1,2,3,4,5), dtype=float32)
    print(tensor)
    # Set to my_pe + 1 so that it doesn't ever show 0.
    tensor[:] = nvshmem.core.my_pe() + 1
    print(tensor)
    # Free Buffer
    nvshmem.core.free_tensor(tensor)
    print("Ending test Torch interop buffer")

def test_interop_cupy():
    print("Testing CuPy interop mem")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy Interop test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    dev = Device(local_rank_per_node)
    array = nvshmem.core.array((1,2,3,4), dtype="float32")
    print(array)
    # Set to my_pe + 1 so that it doesn't ever show 0.
    array[:] = nvshmem.core.my_pe() + 1
    print(array)
    # Free Buffer
    nvshmem.core.free_array(array)
    print("Done testing CuPy Interop mem")

def test_peer_array():
    print("Testing peer array")
    # TODO: TEAM_NODE not showing up because of anonymous enum
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    arr = nvshmem.core.array((256, 1), dtype="float32")

    # PE0 calls get_peer_buffer on the Buffer
    if local_rank_per_node == 0:
        try:
            peer_arr = nvshmem.core.get_peer_array(arr, ((nvshmem.core.my_pe() + 1) % nvshmem.core.n_pes()))
        except Exception as e:
            peer_arr = None
            # If >1 node (NVL domain) exists, this is an error condition. TODO enhance this to differentiate
            print("peer_buf array failed to create. Check to make sure it failed because of >1 node")
        print(peer_arr, arr)
        # Don't need to call free on a peer buf.
        # However, it is safe to do so. nvshmem.core.free knows to skip nvshmem_free()
        #nvshmem.core.free(peer_buf)

    nvshmem.core.free_array(arr)
    print("End peer array test")

def test_peer_tensor():
    print("Testing peer tensor")
    # TODO: TEAM_NODE not showing up because of anonymous enum
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    tensor = nvshmem.core.tensor((256, 1), dtype=float32)

    # PE0 calls get_peer_buffer on the Buffer
    if local_rank_per_node == 0:
        try:
            peer_tensor = nvshmem.core.get_peer_tensor(tensor, ((nvshmem.core.my_pe() + 1) % nvshmem.core.n_pes()))
        except Exception as e:
            peer_tensor = None
            # If >1 node (NVL domain) exists, this is an error condition. TODO enhance this to differentiate
            print("peer_buf tensor failed to create. Check to make sure it failed because of >1 node")
        print(peer_tensor, tensor)
        # Don't need to call free on a peer buf.
        # However, it is safe to do so. nvshmem.core.free knows to skip nvshmem_free()
        #nvshmem.core.free(peer_buf)
    nvshmem.core.free_tensor(tensor)
    print("End peer tensor test")

def test_del_buffer():
    print("Testing buffer scope")
    # Test that when we call del on a buffer, it doesn't actually go away
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    dev = Device(local_rank_per_node)
    buf = nvshmem.core.buffer(1024)
    ptr = buf._mnff.ptr
    try:
        del buf
    except:
        print("Caught early-warning about explicit free")

    print(_mr_references[local_rank_per_node]._mem_references[ptr])
    assert _mr_references[local_rank_per_node]._mem_references[ptr]["ref_count"] > 0
    # This is technically a user error. They can't free the buf if they del it.
    # NVSHMEM Fini-time checking will save them from themself.
    # In the test, I am hacking it to follow the rules
    _mr_references[local_rank_per_node]._mem_references[ptr]["freed"] = True
    _mr_references[local_rank_per_node]._mem_references[ptr]["buffer"].close()
    print("done testing buffer scope")

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--init-type", "-i", type=str, help="Init type to use", choices=["mpi", "uid"], default="uid")
    args = parser.parse_args()
    if args.init_type == "uid":
        uid_init()
    elif args.init_type == "mpi":
        mpi_init()

    test_buffer()
    test_peer_buffer()
    test_peer_array()
    test_peer_tensor()
    test_interop_cupy()
    test_interop_torch()
    test_del_buffer()

    nvshmem.core.finalize()
