#include "tile/init/on_init.h"

#include "tile/base/never_destroyed.h"
#include "tile/base/random.h"
#include "tile/base/thread/mutex.h"
#include "tile/base/thread/scoped_lock.h"

#include <algorithm>
#include <atomic>
#include <map>
#include <mutex>
#include <vector>

namespace tile {
namespace {
std::vector<std::function<void()>> initializer_registry;
std::vector<std::function<void()>> finalizer_registry;

struct AtExitCallbackRegistry {
    Mutex mutex;
    std::vector<std::function<void()>> registry;
};

AtExitCallbackRegistry *
GetAtExitCallbackRegistry()
{
    static NeverDestroyed<AtExitCallbackRegistry> registry;
    return registry.Get();
}

}// namespace
}// namespace tile

namespace tile {
namespace internal {

void
SetAtExitCallback(std::function<void()> callback)
{
    auto &&registry = GetAtExitCallbackRegistry();
    ScopedLock _(registry->mutex);
    registry->registry.push_back(callback);
}

}// namespace internal
}// namespace tile

namespace tile {
namespace detail {

struct Bean {
    std::function<void()> initializer;
    std::function<void()> finalizer;
};

std::atomic<bool> registry_prepared{false};

std::map<std::int32_t, std::vector<Bean>> *
GetStagingRegistry()
{
    static std::map<std::int32_t, std::vector<Bean>> registry;
    return &registry;
}

void
PrepareForRunningCallbacks()
{
    auto &&temp_registry = *GetStagingRegistry();
    auto &&init_registry = initializer_registry;
    auto &&fini_registry = finalizer_registry;

    for (auto &&item : temp_registry) {
        auto &&k = item.first;
        auto &&v = item.second;
        std::shuffle(v.begin(), v.end(), random::Engine());
        for (auto &&f : v) {
            init_registry.push_back(std::move(f.initializer));
            if (f.finalizer) { fini_registry.push_back(std::move(f.finalizer)); }
        }
    }

    std::reverse(fini_registry.begin(), fini_registry.end());
    registry_prepared.store(true, std::memory_order_relaxed);
}

void
RegisterOnInitCallback(std::int32_t priority, std::function<void()> init, std::function<void()> fini)
{
    TILE_CHECK(!registry_prepared.load(std::memory_order_relaxed),
               "Callbacks may only be registered before `tile::Start()` is called");
    auto &&registry = *GetStagingRegistry();
    TILE_CHECK(init != nullptr, "Initializer must be provided");
    registry[priority].push_back({std::move(init), std::move(fini)});
}

void
RunAllInitializers()
{
    PrepareForRunningCallbacks();
    auto &&registry = initializer_registry;

    for (auto &&c : registry) { c(); }
    registry.clear();
}

void
RunAllFinalizers()
{
    auto &&registry = finalizer_registry;
    for (auto &&c : registry) { c(); }
    registry.clear();
    for (auto &&e : GetAtExitCallbackRegistry()->registry) { e(); }
}

}// namespace detail
}// namespace tile
