#include "erl_nif.h"
#include <stdatomic.h>

atomic_ullong LockSlot[1048576];

ERL_NIF_TERM atomTrue;
ERL_NIF_TERM atomFalse;
ERL_NIF_TERM atomUndefined;

int nifLoad(ErlNifEnv *env, void **priv_data, ERL_NIF_TERM){
    atomTrue = enif_make_atom(env, "true");
    atomFalse = enif_make_atom(env, "false");
    atomUndefined = enif_make_atom(env, "undefined");
    return 0;
}

ERL_NIF_TERM tryLock(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){
    unsigned int KeyIx;
    if (!enif_get_uint(env, argv[0], &KeyIx)){
        return enif_make_badarg(env);
    }
    ErlNifPid ThePid;
    enif_self(env, &ThePid);
    atomic_ullong Expected = 0;
    atomic_ullong Val = (atomic_ullong)(ThePid.pid);

    if (atomic_compare_exchange_strong(&(LockSlot[KeyIx]), &Expected, Val)){
        return atomTrue;
    }else{
        ThePid.pid = (ERL_NIF_TERM)Expected;
        if (enif_is_process_alive(env, &ThePid)){
            return atomFalse;
        }else{
            if (atomic_compare_exchange_strong(&LockSlot[KeyIx], &Expected, Val)) {
                return atomTrue;
            }else{
                return atomFalse;
            }
        }
    }
}

ERL_NIF_TERM releaseLock(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){
    unsigned int KeyIx;
    if (!enif_get_uint(env, argv[0], &KeyIx)){
        return enif_make_badarg(env);
    }
    ErlNifPid ThePid;
    enif_self(env, &ThePid);
    atomic_ullong Expected = (atomic_ullong)(ThePid.pid);
    atomic_ullong Val = 0;

    if (atomic_compare_exchange_strong(&LockSlot[KeyIx], &Expected, Val)){
        return atomTrue;
    }else{
        return atomFalse;
    }
}

ERL_NIF_TERM getLockPid(ErlNifEnv *env, int, const ERL_NIF_TERM argv[]){
    unsigned int KeyIx;
    if (!enif_get_uint(env, argv[0], &KeyIx)){
        return enif_make_badarg(env);
    }
    ErlNifPid ThePid;

    atomic_ullong Var = atomic_load(&(LockSlot[KeyIx]));
    if (Var > 0){
        ThePid.pid = (ERL_NIF_TERM)Var;
        return enif_make_pid(env, &ThePid);
    }else{
        return atomUndefined;
    }
}

static ErlNifFunc nifFuncs[] =
    {
        {"tryLock", 1, tryLock},
        {"releaseLock", 1, releaseLock},
        {"getLockPid", 1, getLockPid}
};

ERL_NIF_INIT(eNifLock, nifFuncs, &nifLoad, NULL, NULL, NULL)