#pragma once

#if !defined(MIJIN_ASYNC_SIGNAL_HPP_INCLUDED)
#define MIJIN_ASYNC_SIGNAL_HPP_INCLUDED 1

#include <algorithm>
#include <cstdint>
#include <functional>
#include <memory>
#include <mutex>
#include <vector>
#include "../internal/common.hpp"
#include "../util/flag.hpp"

namespace mijin
{

//
// public defines
//

//
// public constants
//

using signal_token_t = std::uint32_t;
inline constexpr signal_token_t INVALID_SIGNAL_TOKEN = std::numeric_limits<signal_token_t>::max();

//
// public types
//

MIJIN_DEFINE_FLAG(Oneshot);

template<template<typename> typename TAllocator, typename... TArgs>
class BaseSignal
{
public:
    using handler_t = std::function<void(TArgs...)>; // TODO: write a custom function wrapper with allocator support
    using token_t = signal_token_t;
private:
    struct RegisteredHandler
    {
        handler_t callable;
        std::weak_ptr<void> referenced;
        token_t token;
        Oneshot oneshot = Oneshot::NO;
    };
    using handler_vector_t = std::vector<RegisteredHandler, TAllocator<RegisteredHandler>>;
private:
    handler_vector_t handlers_;
    token_t nextToken = 1;
    std::mutex handlersMutex_;
public:
    explicit BaseSignal(TAllocator<void> allocator = {}) : handlers_(TAllocator<RegisteredHandler>(std::move(allocator))) {}
    BaseSignal(const BaseSignal&) = delete;
    BaseSignal(BaseSignal&&) MIJIN_NOEXCEPT = default;

    void reinit(TAllocator<void> allocator = {})
    {
        MIJIN_ASSERT(handlers_.empty(), "Attempting to re-initialize a signal that already has handlers.");
        handlers_ = handler_vector_t(TAllocator<RegisteredHandler>(std::move(allocator)));
    }
public:
    BaseSignal& operator=(const BaseSignal&) = delete;
    BaseSignal& operator=(BaseSignal&&) MIJIN_NOEXCEPT = default;
public:
    template<typename THandler, typename TWeak = void>
    inline token_t connect(THandler handler, Oneshot oneshot = Oneshot::NO, std::weak_ptr<TWeak> referenced = std::weak_ptr<TWeak>()) MIJIN_NOEXCEPT;
    template<typename TObject, typename TWeak = void>
    inline token_t connect(TObject& object, void (TObject::* handler)(TArgs...), Oneshot oneshot = Oneshot::NO, std::weak_ptr<TWeak> referenced = std::weak_ptr<TWeak>()) MIJIN_NOEXCEPT;
    inline void disconnect(token_t token) MIJIN_NOEXCEPT;

    template<typename... TArgs2>
    inline void emit(TArgs2&&... args) MIJIN_NOEXCEPT;
};

template<typename... TArgs>
using Signal = BaseSignal<MIJIN_DEFAULT_ALLOCATOR, TArgs...>;

//
// public functions
//

template<template<typename> typename TAllocator, typename... TArgs>
template<typename THandler, typename TWeak>
inline auto BaseSignal<TAllocator, TArgs...>::connect(THandler handler, Oneshot oneshot, std::weak_ptr<TWeak> referenced) MIJIN_NOEXCEPT -> token_t
{
    std::lock_guard lock(handlersMutex_);

    auto callable = handler_t(handler);
    handlers_.push_back({
        .callable = std::move(callable),
        .referenced = std::move(referenced),
        .token = nextToken,
        .oneshot = oneshot
    });

    return nextToken++;
}

template<template<typename> typename TAllocator, typename... TArgs>
template<typename TObject, typename TWeak>
inline auto BaseSignal<TAllocator, TArgs...>::connect(TObject& object, void (TObject::* handler)(TArgs...), Oneshot oneshot, std::weak_ptr<TWeak> referenced) MIJIN_NOEXCEPT -> token_t
{
    std::lock_guard lock(handlersMutex_);

    auto callable = [object = &object, handler](TArgs... args)
    {
        std::invoke(handler, object, std::forward<TArgs>(args)...);
    };
    handlers_.push_back({
        .callable = std::move(callable),
        .referenced = std::move(referenced),
        .token = nextToken,
        .oneshot = oneshot
    });

    return nextToken++;
}

template<template<typename> typename TAllocator, typename... TArgs>
inline void BaseSignal<TAllocator, TArgs...>::disconnect(token_t token) MIJIN_NOEXCEPT
{
    std::lock_guard lock(handlersMutex_);

    auto it = std::remove_if(handlers_.begin(), handlers_.end(), [token](const RegisteredHandler& handler)
    {
        return handler.token == token;
    });
    handlers_.erase(it, handlers_.end());
}

template<template<typename> typename TAllocator, typename... TArgs>
template<typename... TArgs2>
inline void BaseSignal<TAllocator, TArgs...>::emit(TArgs2&&... args) MIJIN_NOEXCEPT
{
    std::lock_guard lock(handlersMutex_);
    
    // first erase any handlers with expired references
    // auto it = std::remove_if(handlers_.begin(), handlers_.end(), [](const RegisteredHandler& handler)
    // {
    //     return handler.referenced.expired();
    // });
    // handlers_.erase(it, handlers_.end());
    // TODO: this doesn't really work since expired() also returns true if the pointer was never set

    // invoke all handlers
    for (RegisteredHandler& handler : handlers_)
    {
        handler.callable(std::forward<TArgs2>(args)...);
    }

    // remove any oneshot
    auto it = std::remove_if(handlers_.begin(), handlers_.end(), [](const RegisteredHandler& handler)
    {
        return handler.oneshot;
    });
    handlers_.erase(it, handlers_.end());
}

} // namespace mijin

#endif // !defined(MIJIN_ASYNC_SIGNAL_HPP_INCLUDED)