/**
 * @file
 * Reference counting subroutines.
 */
#include "refcount/refcount.h"

#include "refcount/allocator.h"

#include <stdbool.h>
#include <stdlib.h>

/**
 * Callback to be executed when a program error occurred. A sane to default is
 * to print a message and then abort. This is not used to report normal
 * recoverable errors, but rather to report fatal logic errors.
 */
refcount_debug_callback_t refcount_debug_breakpoint = NULL;

/**
 * Call the breakpoint handler if one is installed.
 * @param oper The failing operation
 * @param ctx The #RefcountContext
 * @param obj The object on which the operation tried to act
 */
#define BREAKPOINT(oper, ctx, obj)                                     \
    if (refcount_debug_breakpoint) {                                   \
        refcount_debug_breakpoint(REFCOUNT_OPER_##oper, (ctx), (obj)); \
    }

/**
 * Default context to use for functions that do not take a context. You must
 * initialize this before use.
 */
RefcountContext *refcount_default_context = NULL;

/**
 * Create a new context.
 * @param entry_offset The offset to the #RefcountEntry member
 * @param held_refs_callback A function that will list all references held by an
 * object
 * @param destroy_callback A function to be called when an objects reference
 * count drops to zero
 * @param user_data Extra data to pass to the callbacks
 * @param alloc The #RefcountAllocator to use for all internal allocations
 * @return The new context, or NULL in the case of an error
 */
RefcountContext *
refcount_make_context(size_t entry_offset,
                      refcount_held_refs_callback_t held_refs_callback,
                      refcount_destroy_callback_t destroy_callback,
                      void *user_data, const RefcountAllocator *alloc) {
    if (!alloc) {
        alloc = refcount_global_allocator;
    }
    RefcountContext *ctx = refcount_malloc(alloc, sizeof(RefcountContext));
    if (!ctx) {
        return NULL;
    }
    ctx->entry_offset = entry_offset;
    ctx->held_refs_callback = held_refs_callback;
    ctx->destroy_callback = destroy_callback;
    ctx->user_data = user_data;
    ctx->alloc = *alloc;
    refcount_allocator_to_ht_allocator(&ctx->alloc, &ctx->ht_alloc);

    ctx->static_objects = NULL;
    ctx->gc_roots = NULL;
    return ctx;
}

/**
 * Cleanup a #RefcountContext and free any associated resources. This also
 * causes all remaining references, both static and non-static, to be dropped
 * and the context's destroy callback to be called on each object.
 * @param ctx The #RefcountContext
 */
void refcount_context_destroy(RefcountContext *ctx) {
    refcount_list_free_with_data_full(
        ctx->static_objects, refcount_context_deinit_static_as_callback, ctx,
        &ctx->alloc);

    refcount_context_garbage_collect(ctx);

    refcount_free(&ctx->alloc, ctx);
}

/**
 * Internal magic macro.
 */
#define ENTRY (REFCOUNT_OBJECT_ENTRY(ctx, obj))

/**
 * Initialize the #RefcountEntry member of an object.
 * @param ctx The #RefcountContext
 * @param obj The object to initialize
 */
void refcount_context_init_obj(const RefcountContext *ctx, void *obj) {
    if (obj) {
        ENTRY->is_static = false;
        ENTRY->impl.counted.gc_root = NULL;
        ENTRY->impl.counted.ref_count = 0;
        ENTRY->weak_refs = NULL;
    }
}

/**
 * Register a static object in a context.
 * @param ctx The #RefcountContext
 * @param obj The object to register
 * @return True on success, false otherwise
 */
bool refcount_context_init_static(RefcountContext *ctx, void *obj) {
    ENTRY->is_static = true;
    ctx->static_objects =
        refcount_list_push_full(ctx->static_objects, obj, &ctx->alloc);
    if (!ctx->static_objects) {
        return false;
    }
    ENTRY->impl.static_entry = ctx->static_objects;
    ENTRY->weak_refs = NULL;
    return true;
}

/**
 * Return the references held by an object.
 * @param ctx The #RefcountContext
 * @param obj The object
 * @param refs Where to store the refs
 * @return True on success
 */
static inline bool obj_held_refs(RefcountContext *ctx, void *obj,
                                 RefcountList **refs) {
    if (ctx->held_refs_callback) {
        return ctx->held_refs_callback(obj, refs, ctx->user_data);
    }
    return true;
}

/**
 * Call back used from #clear_weak_references.
 * @param wr_raw The #RefcountWeakref object
 */
static void clear_weak_reference_callback(void *wr_raw) {
    RefcountWeakref *wr = wr_raw;
    wr->entry = NULL;
    wr->data = NULL;
}

/**
 * Clear all weak references of an object without freeing them.
 * @param ctx The #RefcountContext
 * @param obj The object to clear
 */
static inline void clear_weak_references(const RefcountContext *ctx,
                                         void *obj) {
    refcount_list_free_full(ENTRY->weak_refs, clear_weak_reference_callback,
                            &ctx->alloc);
}

/**
 * Unregister a static object in a context.
 * @param ctx The #RefcountContext
 * @param obj The object to unregister
 * @return True on success, false otherwise
 */
bool refcount_context_deinit_static(RefcountContext *ctx, void *obj) {
    if (refcount_context_is_static(ctx, obj)) {
        RefcountList *held_refs = NULL;
        if (!obj_held_refs(ctx, obj, &held_refs)) {
            return false;
        }
        clear_weak_references(ctx, obj);
        ctx->static_objects = refcount_list_remove_full(
            ctx->static_objects, ENTRY->impl.static_entry, NULL, &ctx->alloc);
        refcount_list_free_with_data_full(
            held_refs, refcount_context_unref_as_callback, ctx, &ctx->alloc);
        return true;
    } else {
        BREAKPOINT(DEINIT_STATIC, ctx, obj)
        return false;
    }
}

/**
 * Increment the reference count of an object.
 * @param ctx The #RefcountContext
 * @param obj The object to reference
 * @return The input object
 */
void *refcount_context_ref(const RefcountContext *ctx, void *obj) {
    if (!obj) {
        return NULL;
    } else if (!ENTRY->is_static) {
        ++ENTRY->impl.counted.ref_count;
    }
    return obj;
}

/**
 * Track an object as a GC root in a context. It is safe to call this on an
 * already tracked object.
 * @param ctx the #RefcountContext
 * @param obj The object to track
 * @return True on success
 */
static bool track_gc_root(RefcountContext *ctx, void *obj) {
    if (!ENTRY->impl.counted.gc_root) {
        ctx->gc_roots =
            refcount_list_push_full(ctx->gc_roots, obj, &ctx->alloc);
        if (!ctx->gc_roots) {
            return false;
        }
        ENTRY->impl.counted.gc_root = ctx->gc_roots;
    }
    return true;
}

/**
 * Remove an object from the GV root list of a context. It is safe to call this
 * on an object that is not currently tracked.
 * @param ctx The #RefcountContext
 * @param obj The object to untrack
 */
static void remove_gc_root(RefcountContext *ctx, void *obj) {
    ctx->gc_roots = refcount_list_remove_full(
        ctx->gc_roots, ENTRY->impl.counted.gc_root, NULL, &ctx->alloc);
    ENTRY->impl.counted.gc_root = NULL;
}

/**
 * Decrement the reference count of an object. If the reference count dropped to
 * zero, add it to the queue. It is safe to call this on a static object.
 * @param ctx The #RefcountContext
 * @param obj The object to unref
 * @param queue A double pointer to a #RefcountList acting as a queue
 * @param toplevel The toplevel object that triggered this unref. It is ignored
 * so that it is not freed twice in the specific case that an object with a
 * floating ref and a reference cycle is freed
 * @return NULL if the reference count fell to 0, the given object otherwise
 */
static void *unref_to_queue(RefcountContext *ctx, void *obj,
                            RefcountList **queue, void *toplevel) {
    if (!obj) {
        return NULL;
    } else if (ENTRY->is_static) {
        return obj;
    } else if (obj == toplevel) {
        return NULL;
    } else if (ENTRY->impl.counted.ref_count <= 1) {
        *queue = refcount_list_push_full(*queue, obj, &ctx->alloc);
        return NULL;
    }
    --ENTRY->impl.counted.ref_count;
    track_gc_root(ctx, obj);
    return obj;
}

/**
 * A pair of a context and double pointer to a queue. This is for internal use.
 * @see unref_to_queue_as_callback
 */
struct ContextAndQueue {
    RefcountContext *ctx; //!< The context.
    void *toplevel; //!< Toplevel object that triggered the unref.
    RefcountList **queue; //!< The queue.
};

/**
 * Unref an object using a combined context and queue parameter.
 * @param obj The object
 * @param ctx_and_queue_raw The context and queue
 */
static void unref_to_queue_as_callback(void *obj, void *ctx_and_queue_raw) {
    struct ContextAndQueue *ctx_and_queue = ctx_and_queue_raw;
    unref_to_queue(ctx_and_queue->ctx, obj, ctx_and_queue->queue,
                   ctx_and_queue->toplevel);
}

/**
 * Destroy an object by calling it's destructor.
 * @param ctx The #RefcountContext
 * @param obj The object to destroy
 */
static inline void destroy_object(RefcountContext *ctx, void *obj) {
    clear_weak_references(ctx, obj);
    remove_gc_root(ctx, obj);
    if (ctx->destroy_callback) {
        ctx->destroy_callback(obj, ctx->user_data);
    }
}

/**
 * Continually release held references and object held in a queue.
 * @param ctx The #RefcountContext
 * @param queue The queue
 * @param toplevel Toplevel object that triggered the unref
 */
static void process_unref_queue(RefcountContext *ctx, RefcountList *queue,
                                void *toplevel) {
    struct ContextAndQueue ctx_and_queue = {
        .ctx = ctx, .queue = &queue, .toplevel = toplevel};
    while (queue) {
        void *cur = refcount_list_peek(queue);
        RefcountList *held_refs = NULL;
        queue = refcount_list_pop_full(queue, NULL, &ctx->alloc);
        if (!cur) {
            continue;
        }
        if (obj_held_refs(ctx, cur, &held_refs)) {
            // I don't really know how else to handle this as I can't think of a
            // good way to undo all the unrefs that have already been processed,
            // so we can't really make this atomic without going over all
            // objects twice.
            refcount_list_free_with_data_full(held_refs,
                                              unref_to_queue_as_callback,
                                              &ctx_and_queue, &ctx->alloc);
        }
        destroy_object(ctx, cur);
    }
}

/**
 * Decrement the reference count of a object. It is safe to call this on a
 * static object.
 * @param ctx The #RefcountContext
 * @param obj The object
 * @return NULL if the object's reference counter fell to 0, otherwise the
 * object
 */
void *refcount_context_unref(RefcountContext *ctx, void *obj) {
    if (!obj) {
        return NULL;
    }
    RefcountList *queue = NULL;
    void *retval = unref_to_queue(ctx, obj, &queue, NULL);
    process_unref_queue(ctx, queue, obj);
    return retval;
}

/**
 * Float a reference of an object. That is, decrement it's refrerence count, but
 * don't free it if the count drops to 0. It is safe to call this on a static
 * object. It is an error to call this on an object that already has a floating
 * reference.
 * @param ctx The #RefcountContext
 * @param obj The object
 * @return The input object
 */
void *refcount_context_float(RefcountContext *ctx, void *obj) {
    if (!obj) {
        return NULL;
    } else if (ENTRY->is_static) {
        return obj;
    } else if (!ENTRY->impl.counted.ref_count) {
        BREAKPOINT(FLOAT, ctx, obj);
        return obj;
    }
    if (--ENTRY->impl.counted.ref_count) {
        track_gc_root(ctx, obj);
    } else {
        remove_gc_root(ctx, obj);
    }
    return obj;
}

/**
 * Set of hash table functions used in #check_gc_root.
 */
static const HTTableFunctions ROOT_COUNTS_FNS = {
    .hash = ht_intptr_hash_callback,
    .equal = ht_intptr_equal_callback,
    .destroy_key = NULL,
    .destroy_value = NULL,
    .user_data = NULL,
};

/**
 * Holds data for #free_roots_foreach, used in #check_gc_root.
 */
struct ContextAndRootPtr {
    RefcountContext *ctx; //!< The context.
    RefcountList **root_ptr; //!< Double pointer to the root.
    bool did_update; //!< Weather or not *root_ptr was changed.
};

/**
 * Foreach function to free roots from a hash table. Used in #check_gc_root.
 * @param obj The object to free
 * @param ignored Ignored
 * @param user_data A #ContextAndRootPtr
 * @return Always false
 */
static bool free_roots_foreach(void *obj, void *ignored, void *user_data) {
    if (!obj) {
        return false;
    }
    struct ContextAndRootPtr *data = user_data;
    if (*data->root_ptr
        && *data->root_ptr
               == REFCOUNT_OBJECT_ENTRY(data->ctx, obj)->impl.counted.gc_root) {
        *data->root_ptr = (*data->root_ptr)->next;
        data->did_update = true;
    }
    destroy_object(data->ctx, obj);
    return false;
}

/**
 * Check the root pointed to by the double pointer root_ptr. After the call,
 * root_ptr is set to the next root to be checked.
 * @param ctx The context
 * @param root_ptr Double pointer to one of the GC roots
 * @return The number of object's freed, or -1 if an error happened
 */
static ptrdiff_t check_gc_root(RefcountContext *ctx, RefcountList **root_ptr) {
    HTTable *counts = ht_new(&ROOT_COUNTS_FNS, &ctx->ht_alloc, NULL);
    if (!counts) {
        *root_ptr = (*root_ptr)->next;
        return -1;
    }
    RefcountList *root = *root_ptr;
    RefcountList *queue = NULL;
    if (!obj_held_refs(ctx, root->data, &queue)) {
        ht_free(counts);
        *root_ptr = (*root_ptr)->next;
        return -1;
    }
    size_t seen_objects = 0;
    size_t clear_objects = 0;
    // ignore allocation errors until I decide how to deal with them (in the far
    // future)
    while (queue) {
        void *obj = queue->data;
        if (!obj || refcount_context_is_static(ctx, obj)) {
            queue = refcount_list_pop_full(queue, NULL, &ctx->alloc);
            continue;
        }
        uintptr_t count;
        if (ht_has(counts, obj)) {
            count = HT_UUNSTUFF(ht_get(counts, obj));
        } else {
            count = REFCOUNT_OBJECT_ENTRY(ctx, obj)->impl.counted.ref_count;
            ++seen_objects;
        }
        queue = refcount_list_pop_full(queue, NULL, &ctx->alloc);
        if (count > 0) {
            ht_insert(counts, obj, HT_STUFF(--count));
            if (count == 0) {
                ++clear_objects;
            }
            obj_held_refs(ctx, obj, &queue);
        }
    }
    if (seen_objects == clear_objects) {
        struct ContextAndRootPtr data = {
            .ctx = ctx, .root_ptr = root_ptr, .did_update = false};
        ht_foreach(counts, free_roots_foreach, &data);
        if (!data.did_update) {
            *root_ptr = (*root_ptr)->next;
        }
    } else {
        *root_ptr = (*root_ptr)->next;
    }
    ht_free(counts);
    return seen_objects == clear_objects ? seen_objects : 0;
}

/**
 * Run the garbage collector on a context.
 * @param ctx The #RefcountContext
 * @return The number of object's freed, or -1 if an error occurred
 */
ptrdiff_t refcount_context_garbage_collect(RefcountContext *ctx) {
    if (!ctx->held_refs_callback) {
        // no loops possible
        return 0;
    }
    ptrdiff_t total_cleared = 0;
    RefcountList *root = ctx->gc_roots;
    while (root) {
        ptrdiff_t res = check_gc_root(ctx, &root);
        if (res < 0) {
            return -1;
        }
        total_cleared += res;
    }
    return total_cleared;
}

/**
 * Create a new weak reference for an object. A weak reference will allow safe
 * access to the referenced object without holding a reference. That is, the
 * referenced object can be accessed until it's reference count falls to 0 and
 * it is freed. After this, attempts to use the weak reference will just return
 * NULL to indicate that the referenced object is no longer in existence.
 * @param ctx The #RefcountContext
 * @param obj The object for which to create a weak reference
 * @return The newly created weak reference, or NULL if an allocated error
 * occurred.
 */
RefcountWeakref *refcount_context_make_weakref(const RefcountContext *ctx,
                                               void *obj) {
    RefcountWeakref *wr = refcount_malloc(&ctx->alloc, sizeof(RefcountWeakref));
    if (!wr) {
        return NULL;
    }
    RefcountList *new_weak_refs =
        refcount_list_push_full(ENTRY->weak_refs, wr, &ctx->alloc);
    if (!new_weak_refs) {
        refcount_free(&ctx->alloc, wr);
        return NULL;
    }
    ENTRY->weak_refs = new_weak_refs;
    wr->entry = new_weak_refs;
    wr->data = obj;
    return wr;
}

/**
 * Destroy a weak reference. This has no effect on the reference count of the
 * original object.
 * @param ctx The #RefcountContext
 * @param wr The weak reference
 */
void refcount_context_destroy_weakref(const RefcountContext *ctx,
                                      RefcountWeakref *wr) {
    if (wr->data) {
        void *obj = wr->data;
        ENTRY->weak_refs = refcount_list_remove_full(
            ENTRY->weak_refs, wr->entry, NULL, &ctx->alloc);
    }
    refcount_free(&ctx->alloc, wr);
}

/**
 * Add a reference to an object referenced by a weak reference and return the
 * object. If the referenced object no longer exists, return NULL.
 * @param ctx The #RefcountContext
 * @param wr The weak reference
 * @return The newly referenced object, or NULL
 */
void *refcount_context_ref_weakref(const RefcountContext *ctx,
                                   const RefcountWeakref *wr) {
    if (!wr->data) {
        return NULL;
    }
    return refcount_context_ref(ctx, wr->data);
}