#include "cbag.h"

// shared variables
// block_t *global_head_block[NR_THREADS];
// initialize to num_threads (once known)
block_t **global_head_block;
// thread-local storage
block_t *thread_block;
_Atomic(block_t *) steal_prev;
block_t *steal_block;
bool found_add;
int thread_head, steal_head, steal_index;
int num_threads;
int thread_id;

#pragma omp threadprivate(thread_block, steal_block, steal_prev, found_add, thread_head, steal_head, steal_index, thread_id)

block_t *block_get_ref(block_t *block)
{
    return (block_t *)(((long long unsigned int)block) & MASK_POINTER);
}

bool block_get_mark1(block_t *block)
{
    return (((long long unsigned int)block) & MASK_MARK1) > 0;
}

bool block_get_mark2(block_t *block)
{
    return (((long long unsigned int)block) & MASK_MARK2) > 0;
}

void block_set_mark1(block_t **block, bool mark1)
{
    *block = (block_t *)((((long long unsigned int)(*block)) & (MASK_POINTER + MASK_MARK2)) + (mark1 ? MASK_MARK1 : 0));
}

void block_set_mark2(block_t **block, bool mark2)
{
    *block = (block_t *)((((long long unsigned int)(*block)) & (MASK_POINTER + MASK_MARK1)) + (mark2 ? MASK_MARK2 : 0));
}

void mark_1_block(block_t *block)
{
    for (;;)
    {
        block_t *next = block->next;
        // replacement is {next.p, next.mark2, true};
        block_t *replacement = next;
        block_set_mark1(&replacement, true);
        // TODO: write wrapper for marking idea instead of these weird dual struct construction!
        if (block_get_ref(next) == NULL || block_get_mark1(next) || atomic_compare_exchange_strong(&block->next, &next, replacement))
        {
            break;
        }
        // C11 atomic behavior beware: atomic_compare_exchange writes value of &block->next into next upon fail!
    }
}

block_t *new_block()
{
    block_t *block = (block_t *)malloc(sizeof(block_t));
#if VERIFY
    assert(block != NULL);
#endif
    block->next = NULL;
    notify_all(block);
    for (int i = 0; i < BLOCK_SIZE; i++)
        block->nodes[i] = NULL;
    return block;
}

void notify_all(block_t *block)
{
    for (int i = 0; i < num_threads / WORD_SIZE; i++)
        block->notifyAdd[i] = 0;
}

void notify_start(block_t *block, int id)
{
    uint64_t old;
    do
    {
        old = block->notifyAdd[id / WORD_SIZE];
    } while (!atomic_compare_exchange_strong(&(block->notifyAdd[id / WORD_SIZE]), &old, old | (1 << (id % WORD_SIZE))));
}

bool notify_check(block_t *block, int id)
{
    return (block->notifyAdd[id / WORD_SIZE] & (1 << (id % WORD_SIZE))) == 0;
}

void cbag_init_bag(int thread_count)
{
    num_threads = thread_count;
    TRACE(printf("init bag num threads %i\n", num_threads));
    global_head_block = (block_t **)malloc(sizeof(block_t *) * num_threads);
#if VERIFY
    assert(global_head_block != NULL);
#endif
    for (int i = 0; i < num_threads; i++)
        global_head_block[i] = NULL;
}

void cbag_destroy_bag()
{
#if VERIFY
    assert(global_head_block != NULL);
#endif
    free(global_head_block);
}

void cbag_init_thread()
{
    thread_id = omp_get_thread_num();
    thread_block = global_head_block[thread_id];
    thread_head = BLOCK_SIZE;
    steal_index = 0;
    steal_block = NULL;
    steal_prev = NULL;
    steal_head = BLOCK_SIZE;
}

void cbag_add(void *item)
{
    TRACE(printf("thread %2d inserting item %p\n", thread_id, item));
    // TRACE(printf("    %d %d %d %d %d %p %p %p\n", found_add, thread_head, steal_head, steal_index, thread_id, thread_block, steal_block, steal_prev));
    int head = thread_head;
    block_t *block = thread_block;
    for (;;)
    {
        if (head == BLOCK_SIZE)
        {
            block_t *oldblock = block;
            block = new_block();
            block->next = oldblock;
            global_head_block[thread_id] = block;
            thread_block = block;
            TRACE(printf("thread %2d created new block (oldblock %p newblock %p global_head_block[%i] %p\n", thread_id, oldblock, block, thread_id, global_head_block[thread_id]));
            head = 0;
        }
        else if (block->nodes[head] == NULL)
        {
            notify_all(block);
            block->nodes[head] = item;
            thread_head = head + 1;
            TRACE(printf("thread %2d successfully inserted item %p\n", thread_id, item));
            return;
        }
        else
            head++;
    }
}

void *cbag_try_remove_any()
{
    int head = thread_head - 1;
    block_t *block = thread_block;
    int round = 0;
    for (;;)
    {
        // TRACE(printf("block %p, block->next %p, blockref %p, block->nextref %p\n", block, block != NULL ? block->next : NULL, block_get_ref(block), block != NULL ? block_get_ref(block->next) : NULL));
        if (block == NULL || (head < 0 && block_get_ref(block->next) ==
                                              NULL))
        {
            TRACE(printf("thread %2d try_remove_any, attempting to steal block\n", thread_id));
            do
            {
                int i = 0;
                do
                {
                    void *result = try_steal_block(round);
                    TRACE(printf("thread %2d try_remove_any, try_steal_block returned %p, found add %i, i %i, steal_block %p, num_threads %i\n", thread_id, result, found_add, i, steal_block, num_threads));
                    if (result != NULL)
                        return result;
                    if (found_add)
                    {
                        round = 0;
                        i = 0;
                    }
                    else if (steal_block == NULL)
                        i++;
                } while (i < num_threads);
            } while (++round <= num_threads);
            return NULL;
        }
        if (head < 0)
        {
            mark_1_block(block);
            for (;;)
            {
                block_t *next = block->next; // de_ref_link(&block->next);
                if (block_get_mark2(next))
                    mark_1_block(block_get_ref(next));
                if (block_get_mark1(next))
                {
                    if (block_get_ref(next))
                        notify_all(block_get_ref(next));
                    if (atomic_compare_exchange_strong(&(global_head_block[thread_id]),
                                                       &block, block_get_ref(next)))
                    {
                        // {NULL, false, true};
                        block->next = (block_t *)MASK_MARK1;
                        // delete_node(block);
                        // re_scan(next);
                        block = block_get_ref(next);
                    }
                    else
                        block = global_head_block[thread_id]; // de_ref_link(&global_head_block[thread_id]);
                }
                else
                    break;
            }
            thread_block = block;
            thread_head = BLOCK_SIZE;
            head = BLOCK_SIZE - 1;
        }
        else
        {
            void *data = block->nodes[head];
            if (data == NULL)
                head--;
            else if (atomic_compare_exchange_strong(&block->nodes[head], &data, NULL))
            {
                thread_head = head;
                TRACE(printf("thread %2d removed item %p\n", thread_id, data));
                return data;
            }
        }
    }
}

void *try_steal_block(int round)
{
    int head = steal_head;
    block_t *block = steal_block;
    found_add = false;
    TRACE(printf("thread %2d try_steal_block start, steal_index %i, steal_head %i, steal_block %p, steal_prev %p\n", thread_id, steal_index, steal_head, steal_block, steal_prev));
    if (block == NULL)
    {
        block = global_head_block[steal_index]; // de_ref_link(&global_head_block[steal_index]);
        steal_block = block;
        steal_head = head = 0;
    }
    if (head == BLOCK_SIZE)
    {
        steal_block = block = next_steal_block(block);
        head = 0;
    }
    if (block == NULL)
    {
        steal_index = (steal_index + 1) % num_threads;
        steal_head = 0;
        steal_block = NULL;
        steal_prev = NULL;
        TRACE(printf("thread %2d try_steal_block unsuccessful, steal_index %i, steal_head %i, steal_block %p, steal_prev %p\n", thread_id, steal_index, steal_head, steal_block, steal_prev));
        return NULL;
    }
    if (round == 1)
        notify_start(block, thread_id);
    else if (round > 1 && notify_check(block, thread_id))
        found_add = true;
    for (;;)
    {
        if (head == BLOCK_SIZE)
        {
            steal_head = head;
            return NULL;
        }
        else
        {
            void *data = block->nodes[head];
            if (data == NULL)
                head++;
            else if (atomic_compare_exchange_strong(&block->nodes[head], &data, NULL))
            {
                TRACE(printf("thread %2d try_steal_block successfully stole %p, steal_index %i, steal_head %i, steal_block %p, steal_prev %p\n", thread_id, data, steal_index, steal_head, steal_block, steal_prev));
                steal_head = head;
                return data;
            }
        }
    }
}

block_t *next_steal_block(block_t *block)
{
    block_t *next;
    for (;;)
    {
        if (block == NULL)
        {
            block = global_head_block[steal_index]; // de_ref_link(&global_head_block[steal_index]);
            break;
        }
        next = block->next; // de_ref_link(&block->next);
        if (block_get_mark2(next))
            mark_1_block(block_get_ref(next));
        if (steal_prev == NULL || block_get_ref(next) == NULL)
        {
            if (block_get_mark1(next))
            {
                if (block_get_ref(next))
                    notify_all(block_get_ref(next));
                if (atomic_compare_exchange_strong(&global_head_block[steal_index], &block, block_get_ref(next)))
                {
                    // {NULL, false, true};
                    block->next = (block_t *)MASK_MARK1;
                    // delete_node(block);
                    // re_scan(next);
                }
                else
                {
                    steal_prev = NULL;
                    block = global_head_block[steal_index]; // de_ref_link(&global_head_block[steal_index]);
                    continue;
                }
            }
            else
                steal_prev = block;
        }
        else
        {
            if (block_get_mark1(next))
            {
                // {block, steal_prev->next.mark2, false};
                block_t *prevnext = block;
                block_set_mark1(&prevnext, false);
                block_set_mark2(&prevnext, block_get_mark2(steal_prev->next));
                if (atomic_compare_exchange_strong(&steal_prev->next, &prevnext, block_get_ref(next)))
                {
                    // {NULL, false, true};
                    block->next = (block_t *)MASK_MARK1;
                    // delete_node(block);
                    // re_scan(next);
                }
                else
                {
                    steal_prev = NULL;
                    block = global_head_block[steal_index]; // de_ref_link(&global_head_block[steal_index]);
                    continue;
                }
            }
            else if (block == steal_block)
            {
                // {block, true, false};
                block_t *replacement = block;
                block_set_mark1(&replacement, false);
                block_set_mark2(&replacement, true);
                if (atomic_compare_exchange_strong(&steal_prev->next, &block, replacement))
                {
                    mark_1_block(block);
                    continue;
                }
                else
                {
                    steal_prev = NULL;
                    block = global_head_block[steal_index]; // de_ref_link(&global_head_block[steal_index]);
                    continue;
                }
            }
            else
                steal_prev = block;
        }
        if (block == steal_block || block_get_ref(next) == steal_block)
        {
            block = block_get_ref(next);
            break;
        }
        block = block_get_ref(next);
    }
    return block;
}