#include "erl_nif.h"
#include <atomic>
#include <forward_list>
#include <list>
#include <set>
#include <map>
#include <unordered_set>
using namespace std;

const int LockSize = 2097152;
const int ScheduleCnt = 50;
const int HashSalt = 786234121;
atomic<uint64_t> LockSlot[LockSize];

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;
}

bool isNotCurLocked(forward_list<int> *locked, int KeyIx)
{
    for (auto it : *locked)
    {
        if (it == KeyIx)
            return false;
    }
    return true;
}

bool lockOne(ErlNifEnv *env, ErlNifPid *ThePid, int KeyIx, uint64_t Val)
{
    uint64_t Expected = 0;
    if (LockSlot[KeyIx].compare_exchange_strong(Expected, Val))
    {
        return true;
    }
    else
    {
        ThePid->pid = (ERL_NIF_TERM)Expected;
        if (enif_is_process_alive(env, ThePid))
        {
            return false;
        }
        else
        {
            if (LockSlot[KeyIx].compare_exchange_strong(Expected, Val))
            {
                return true;
            }
            else
            {
                return false;
            }
        }
    }
}

ERL_NIF_TERM tryLockSchedule(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    //enif_fprintf(stdout, "IMY************tryLocksSchedule00 %T \n", argv[0]);
    set<int> locked;
    ERL_NIF_TERM allList = argv[1];
    ERL_NIF_TERM head;
    int KeyIx;
    while (enif_get_list_cell(env, allList, &head, &allList))
    {
        enif_get_int(env, head, &KeyIx);
        locked.insert(KeyIx);
    }

    ErlNifPid ThePid;
    enif_self(env, &ThePid);
    uint64_t Val = (uint64_t)(ThePid.pid);
    allList = argv[0];
    //enif_fprintf(stdout, "IMY************tryLocksSchedule222 %T \n", allList);
    auto search = locked.end();
    while (enif_get_list_cell(env, allList, &head, &allList))
    {
        KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, head, HashSalt) % LockSize;
        //enif_fprintf(stdout, "IMY************tryLocksSchedule222000 %d \n", KeyIx);
             search = locked.find(KeyIx);
            if (search != locked.end())
        {
            if (!lockOne(env, &ThePid, KeyIx, Val))
            {
                goto Failed;
            }
            else
            {
                locked.insert(KeyIx);
            }
        }
    }
    //enif_fprintf(stdout, "IMY************tryLocksSchedule333 %T \n", allList);
    return atomTrue;

Failed:
{
    uint64_t RExpected;
    for (auto it : locked)
    {
        RExpected = Val;
        LockSlot[it].compare_exchange_strong(RExpected, 0);
    }
    return atomFalse;
}
}

ERL_NIF_TERM tryLock(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    if (enif_is_list(env, argv[0]))
    {
        // enif_fprintf(stdout, "IMY************tryLocks00 %T \n", argv[0]);
        ERL_NIF_TERM allList = argv[0];
        ERL_NIF_TERM head;
        ErlNifPid ThePid;
        enif_self(env, &ThePid);
        uint64_t Val = (uint64_t)(ThePid.pid);
        int KeyIx;
        //forward_list<int> locked(15);
        //list<int> locked;
        set<int> locked;
       // map<int, int> locked;
        unsigned int cnt = 0;
        auto search = locked.end();
        // enif_fprintf(stdout, "IMY************tryLocks222 %T \n", allList);
        while (enif_get_list_cell(env, allList, &head, &allList))
        {
            // enif_fprintf(stdout, "IMY************tryLocks222---- %T  %T\n", head, allList);
            KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, head, HashSalt) % LockSize;
            // enif_fprintf(stdout, "IMY************tryLocks222000 %d \n", KeyIx);
             search = locked.find(KeyIx);
            if (search != locked.end())
            //if (isNotCurLocked(&locked, KeyIx))
            {
                if (!lockOne(env, &ThePid, KeyIx, Val))
                {
                    goto Failed;
                }
                else
                {
                    //          enif_fprintf(stdout, "IMY************tryLocks222111 %d \n", KeyIx);
                    //locked.push_front(KeyIx);
                    //locked.push_back(KeyIx);
                    locked.insert(KeyIx);
                    //locked.insert({KeyIx, 0});
                    cnt++;
                    //                     enif_fprintf(stdout, "IMY************tryLocks222222 %d \n", cnt);
                    if (cnt >= ScheduleCnt)
                    {
                        goto Schedule;
                    }
                }
            }
        }
        // enif_fprintf(stdout, "IMY************tryLocks333 %T \n", allList);
        return atomTrue;

    Schedule:
    {
        ERL_NIF_TERM ListLocked = enif_make_list(env, 0);
        //for (auto it : locked)
       // {
           // enif_make_list_cell(env, enif_make_int(env, it), ListLocked);
        //}

        ERL_NIF_TERM newargv[2] = {allList, ListLocked};
        //enif_fprintf(stdout, "IMY************tryLocks4444 %T \n", newargv);
        return enif_schedule_nif(env, "tryLockSchedule", 0, &tryLockSchedule, 2, newargv);
    }

    Failed:
    {
        //uint64_t RExpected;
        //for (auto it : locked)
       // {
//RExpected = Val;
        //    LockSlot[it].compare_exchange_strong(RExpected, 0);
       // }
        return atomFalse;
    }
}
else
{
    int KeyIx;
    KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, argv[0], HashSalt) % LockSize;
    ErlNifPid ThePid;
    enif_self(env, &ThePid);
    uint64_t Val = (uint64_t)(ThePid.pid);

    if (lockOne(env, &ThePid, KeyIx, Val))
    {
        return atomTrue;
    }
    else
    {
        return atomFalse;
    }
}
}

ERL_NIF_TERM releaseLockSchedule(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    ERL_NIF_TERM allList = argv[0];
    ERL_NIF_TERM head;
    ErlNifPid ThePid;
    enif_self(env, &ThePid);
    uint64_t Expected = (uint64_t)(ThePid.pid);
    uint64_t RExpected;
    int KeyIx;
    int isAllOk = 1;
    enif_get_int(env, argv[1], &isAllOk);

    while (enif_get_list_cell(env, allList, &head, &allList))
    {
        KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, head, HashSalt) % LockSize;
        RExpected = Expected;
        if (!LockSlot[KeyIx].compare_exchange_strong(RExpected, 0))
        {
            isAllOk = 0;
        }
    }
    return isAllOk > 0 ? atomTrue : atomFalse;
}

ERL_NIF_TERM releaseLock(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    if (enif_is_list(env, argv[0]))
    {

        ERL_NIF_TERM allList = argv[0];
        ERL_NIF_TERM head;
        ErlNifPid ThePid;
        enif_self(env, &ThePid);
        uint64_t Expected = (uint64_t)(ThePid.pid);
        uint64_t RExpected;
        int KeyIx;
        unsigned int cnt = 0;
        int isAllOk = 1;

        while (enif_get_list_cell(env, allList, &head, &allList))
        {
            KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, head, HashSalt) % LockSize;
            RExpected = Expected;
            if (!LockSlot[KeyIx].compare_exchange_strong(RExpected, 0))
            {
                isAllOk = 0;
            }
            cnt++;
            if (cnt >= ScheduleCnt)
            {
                goto Schedule;
            }
        }
        return isAllOk > 0 ? atomTrue : atomFalse;

    Schedule:
    {
        ERL_NIF_TERM ArgvAllOk = enif_make_int(env, isAllOk);
        ERL_NIF_TERM newargv[2] = {allList, ArgvAllOk};
        //enif_fprintf(stdout, "IMY************tryLocks4444 %T \n", newargv);
        return enif_schedule_nif(env, "releaseLockSchedule", ERL_NIF_DIRTY_JOB_CPU_BOUND, &releaseLockSchedule, 2, newargv);
    }
}
else
{
    int KeyIx;
    KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, argv[0], HashSalt) % LockSize;
    ErlNifPid ThePid;
    enif_self(env, &ThePid);
    uint64_t Expected = (uint64_t)(ThePid.pid);

    if (LockSlot[KeyIx].compare_exchange_strong(Expected, 0))
    {
        return atomTrue;
    }
    else
    {
        return atomFalse;
    }
}
}

ERL_NIF_TERM getLockPid(ErlNifEnv *env, int, const ERL_NIF_TERM argv[])
{
    int KeyIx;
    KeyIx = enif_hash(ERL_NIF_INTERNAL_HASH, argv[0], HashSalt) % LockSize;
    ErlNifPid ThePid;

    uint64_t Var = LockSlot[KeyIx].load();
    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)