#pragma once

// Author: Rahil Baber
// Date: 2019/03/17
//
// A custom shared_ptr implementation called stoppable_ptr. Has a threadsafe stop() member function that
// will stop further copies being made (a resetted stoppable_ptr is returned instead). The stop() function
// acts atomically to ensure after it is called the use_count can only decrease. The is_stopped() function
// can be used to query if stop() has been called.
//
// Notable supported features include:
// an aliasing constructor,
// support for array allocation,
// weak_stoppable_ptr,
// enable_stoppable_from_this,
// static_pointer_cast(), dynamic_pointer_cast(), const_pointer_cast(), reinterpret_pointer_cast(),
// make_stoppable(),
// support for custom deleters,
// construction from unique_ptr
//
// Notable differences to shared_ptr:
// stop() and is_stopped() member functions in stoppable_ptr and weak_stoppable_ptr,
// no support for custom allocators (this may be added in the future),
// make_stoppable() will not create the control block and underlying resource in one memory allocation (this
// was purposefully done to make when the resource cleanup happens independent of how the stoppable_ptr was
// created).

#include <mutex>
#include <cstdint>
#include <cstddef>
#include <utility>
#include <memory>
#include <type_traits>

namespace util
{

// Forward declare.
template<typename T>
class stoppable_ptr;

namespace priv_impl_stoppable_ptr
{
class ControlBlockBase
{
public:
    virtual ~ControlBlockBase() = default;

    std::mutex mtx;
    bool stopped;
    std::uint64_t strongRefCount;
    std::uint64_t weakRefCount;

    bool IsStopped() noexcept
    {
        std::lock_guard<std::mutex> lock(mtx);
        return stopped;
    }

    void Stop() noexcept
    {
        std::lock_guard<std::mutex> lock(mtx);
        stopped = true;
    }

    // Returns true if should cleanup the control block.
    virtual bool DecStrongRef() noexcept = 0;

    // Returns true if strongRefCount was increased.
    bool IncStrongRef() noexcept
    {
        std::lock_guard<std::mutex> lock(mtx);
        if (strongRefCount == 0 || stopped)
            return false;

        strongRefCount++;
        return true;
    }

    // Returns true if weakRefCount was increased.
    bool IncWeakRef() noexcept
    {
        std::lock_guard<std::mutex> lock(mtx);
        if (strongRefCount == 0 || stopped)
            return false;

        weakRefCount++;
        return true;
    }

    // Returns true if should cleanup the control block.
    bool DecWeakRef() noexcept
    {
        std::lock_guard<std::mutex> lock(mtx);
        weakRefCount--;

        if (weakRefCount == 0 && strongRefCount == 0)
            return true;

        return false;
    }
};

template<typename T>
class ControlBlockDerived : public ControlBlockBase
{
public:
    T* resource;
    std::function<void(T*)> deleter;

    virtual bool DecStrongRef() noexcept override
    {
        bool delCtrlBlock = false;
        bool delResource  = false;
        T*   resourceCopy;
        {
            std::lock_guard<std::mutex> lock(mtx);
            strongRefCount--;

            if (strongRefCount == 0)
            {
                // Delete the resource outside of the locked scope.
                delResource  = true;
                resourceCopy = resource;
                resource     = nullptr;

                if (weakRefCount == 0)
                    delCtrlBlock = true;
            }
        }

        if (delResource)
            deleter(resourceCopy);

        return delCtrlBlock;
    }
};

template<typename T, typename = void>
class HelperSetWeakThis
{
public:
    static void SetWeakThis(T* ptr, const stoppable_ptr<T>& sp) {}
};

template<typename T>
class HelperSetWeakThis<T, 
    std::enable_if_t<std::is_convertible_v<T, typename T::priv_impl_esft_type>>>
{
public:
    static void SetWeakThis(T* ptr, const stoppable_ptr<T>& sp) { ptr->weak_this = sp; }
};

} // End of namespace priv_impl_stoppable_ptr.

// Forward declare.
template<typename T>
class weak_stoppable_ptr;

template<typename T>
class enable_stoppable_from_this;

template<typename T>
class stoppable_ptr
{
public:
    constexpr stoppable_ptr() noexcept : ptr{nullptr}, ctrl{nullptr} {}
    constexpr stoppable_ptr(std::nullptr_t) noexcept : stoppable_ptr{} {}

    template<typename U>
    explicit stoppable_ptr(U* ptr) 
    { Constructor(ptr, [](T* ptr){ delete ptr; }); }
	
    template<typename U, typename Deleter> 
    stoppable_ptr(U* ptr, Deleter d) 
    { Constructor(ptr, [del=std::move(d)](T* ptr){ del(ptr); }); }

    template<typename Deleter> 
    stoppable_ptr(std::nullptr_t, Deleter d)
    { Constructor(nullptr, [del=std::move(d)](T* ptr){ del(ptr); }); }

    template<typename U> 
    stoppable_ptr(const stoppable_ptr<U>& r, T* ptr) noexcept
    {
        if (r.ctrl && r.ctrl->IncStrongRef())
        {
            ctrl      = r.ctrl;
            this->ptr = ptr;
            return;
        }

        ctrl      = nullptr;
        this->ptr = nullptr;
        return;
    }

    template<typename U> 
    stoppable_ptr(stoppable_ptr<U>&& r, T* ptr) noexcept
    {
        if (r.ctrl)
        {
            ctrl      = r.ctrl;
            this->ptr = ptr;

            r.ctrl = nullptr;
            r.ptr  = nullptr;
            return;
        }

        ctrl      = nullptr;
        this->ptr = nullptr;
        return;
    }
 
    stoppable_ptr(const stoppable_ptr& r) noexcept : stoppable_ptr{r, r.ptr} {}

    template<typename U> 
    stoppable_ptr(const stoppable_ptr<U>& r) noexcept : stoppable_ptr{r, r.ptr} {}

    stoppable_ptr(stoppable_ptr&& r) noexcept : stoppable_ptr{std::move(r), r.ptr} {}

    template<typename U> 
    stoppable_ptr(stoppable_ptr<U>&& r) noexcept : stoppable_ptr{std::move(r), r.ptr} {}

    template<typename U>
    explicit stoppable_ptr(const weak_stoppable_ptr<U>& r) noexcept;

    template<typename U, typename Deleter> 
    stoppable_ptr(std::unique_ptr<U, Deleter>&& r)
    {
        Deleter d = r.get_deleter();
        auto deleter = [del=std::move(d)](T* ptr){ del(ptr); };
        Constructor(r.release(), deleter);
    }

    ~stoppable_ptr() { reset(); }

    stoppable_ptr& operator=(const stoppable_ptr& r) noexcept
    { return operator= <T>(r); }

    template<typename U>
    stoppable_ptr& operator=(const stoppable_ptr<U>& r) noexcept
    {
        stoppable_ptr(r).swap(*this);
        return *this;
    }

    stoppable_ptr& operator=(stoppable_ptr&& r) noexcept
    { return operator= <T>(std::move(r)); }

    template<typename U>
    stoppable_ptr& operator=(stoppable_ptr<U>&& r) noexcept
    {
        stoppable_ptr(std::move(r)).swap(*this);
        return *this;
    }

    template<typename U, typename Deleter> 
    stoppable_ptr& operator=(std::unique_ptr<U, Deleter>&& r)
    {
        stoppable_ptr(std::move(r)).swap(*this);
        return *this;
    }

    T& operator*() const noexcept { return *ptr; }
    T* operator->() const noexcept { return ptr; }
    T& operator[](std::ptrdiff_t idx) const { return ptr[idx]; }
    explicit operator bool() const noexcept { return (ptr != nullptr); }

    void reset() noexcept
    {
        if (ptr)
            ptr = nullptr;

        if (!ctrl)
            return;

        if (ctrl->DecStrongRef())
            delete ctrl;

        ctrl = nullptr;
        return;
    }

    template<typename U>
    void reset(U* ptr) { stoppable_ptr(ptr).swap(*this); }

    template<typename U, typename Deleter> 
    void reset(U* ptr, Deleter d) { stoppable_ptr(ptr, d).swap(*this); }

    T* get() const noexcept { return ptr; }

    std::uint64_t use_count() const noexcept
    {
        if (!ctrl)
            return 0;

        std::lock_guard<std::mutex> lock(ctrl->mtx);
        return ctrl->strongRefCount;
    }

    void swap(stoppable_ptr& r) noexcept
    {
        if (ctrl != r.ctrl)
            std::swap(ctrl, r.ctrl);
       
        if (ptr != r.ptr)
            std::swap(ptr, r.ptr);

        return;
    }

    void stop() const noexcept { if (ctrl) ctrl->Stop(); }
    bool is_stopped() const noexcept { return ctrl ? ctrl->IsStopped() : false; }

private:
    T* ptr;
    priv_impl_stoppable_ptr::ControlBlockBase* ctrl;

    void Constructor(T* ptr, std::function<void(T*)> deleter)
    {
        priv_impl_stoppable_ptr::ControlBlockDerived<T>* derivedCtrl = nullptr;

        try
        {
            derivedCtrl = new priv_impl_stoppable_ptr::ControlBlockDerived<T>();
            ctrl = derivedCtrl;
            derivedCtrl->stopped  = false;
            derivedCtrl->resource = ptr;
            derivedCtrl->deleter  = std::move(deleter);
            derivedCtrl->strongRefCount = 1;
            derivedCtrl->weakRefCount   = 0;

            this->ptr = ptr;
        }
        catch(...)
        {
            deleter(ptr);
            delete derivedCtrl;
            throw;
        }

        try
        {
            priv_impl_stoppable_ptr::HelperSetWeakThis<T>::SetWeakThis(ptr, *this);
        }
        catch(...)
        {
            derivedCtrl->deleter(ptr);
            delete derivedCtrl;
            throw;
        }
    }

    template<typename> friend class stoppable_ptr;
    template<typename> friend class weak_stoppable_ptr;
};

template<typename T, typename U>
bool operator==(const stoppable_ptr<T>& lhs, const stoppable_ptr<U>& rhs) noexcept
{ return (lhs.get()==rhs.get()); }

template<typename T, typename U>
bool operator!=(const stoppable_ptr<T>& lhs, const stoppable_ptr<U>& rhs) noexcept
{ return (lhs.get()!=rhs.get()); }

template<typename T> 
bool operator==(const stoppable_ptr<T>& lhs, std::nullptr_t rhs) noexcept
{ return !lhs.get(); }

template<typename T>
bool operator==(std::nullptr_t lhs, const stoppable_ptr<T>& rhs) noexcept
{ return !rhs.get(); }

template<typename T>
bool operator!=(const stoppable_ptr<T>& lhs, std::nullptr_t rhs) noexcept
{ return lhs.get(); }

template<typename T>
bool operator!=(std::nullptr_t lhs, const stoppable_ptr<T>& rhs) noexcept
{ return rhs.get(); }

template<typename T, typename U> 
stoppable_ptr<T> static_pointer_cast(const stoppable_ptr<U>& r) noexcept
{ return stoppable_ptr<T>(r, static_cast<T*>(r.get())); }

template<typename T, typename U> 
stoppable_ptr<T> static_pointer_cast(stoppable_ptr<U>&& r) noexcept
{ return stoppable_ptr<T>(std::move(r), static_cast<T*>(r.get())); }

template<typename T, typename U> 
stoppable_ptr<T> dynamic_pointer_cast(const stoppable_ptr<U>& r) noexcept
{
    T* p = dynamic_cast<T*>(r.get());
    return p ? stoppable_ptr<T>(r, p) : stoppable_ptr<T>();
}

template<typename T, typename U> 
stoppable_ptr<T> dynamic_pointer_cast(stoppable_ptr<U>&& r) noexcept
{
    T* p = dynamic_cast<T*>(r.get());
    return p ? stoppable_ptr<T>(std::move(r), p) : stoppable_ptr<T>();
}

template<typename T, typename U> 
stoppable_ptr<T> const_pointer_cast(const stoppable_ptr<U>& r) noexcept
{ return stoppable_ptr<T>(r, const_cast<T*>(r.get())); }

template<typename T, typename U> 
stoppable_ptr<T> const_pointer_cast(stoppable_ptr<U>&& r) noexcept
{ return stoppable_ptr<T>(std::move(r), const_cast<T*>(r.get())); }

template<typename T, typename U> 
stoppable_ptr<T> reinterpret_pointer_cast(const stoppable_ptr<U>& r) noexcept
{ return stoppable_ptr<T>(r, reinterpret_cast<T*>(r.get())); }

template<typename T, typename U> 
stoppable_ptr<T> reinterpret_pointer_cast(stoppable_ptr<U>&& r) noexcept
{ return stoppable_ptr<T>(std::move(r), reinterpret_cast<T*>(r.get())); }

namespace priv_impl_stoppable_ptr
{
template<typename T>
class MakeSelector { public: using SingleObjectRetType = stoppable_ptr<T>; };

template<typename T>
class MakeSelector<T[]> { public: using UnknownBoundRetType = stoppable_ptr<T>; };

template<typename T, std::size_t N>
class MakeSelector<T[N]> { public: using KnownBoundRetType = stoppable_ptr<T>; };
} // End of namespace priv_impl_stoppable_ptr.

template<typename T, typename... Args>
typename priv_impl_stoppable_ptr::MakeSelector<T>::SingleObjectRetType
make_stoppable(Args&&... args)
{
    return stoppable_ptr<T>(new T(std::forward<Args>(args)...));
}

template<typename T>
typename priv_impl_stoppable_ptr::MakeSelector<T>::UnknownBoundRetType
make_stoppable(std::size_t N)
{
    using U = typename std::remove_extent<T>::type;
    return stoppable_ptr<U>(new U[N], [](U* ptr){ delete[] ptr; });
}

template<typename T>
typename priv_impl_stoppable_ptr::MakeSelector<T>::KnownBoundRetType
make_stoppable()
{
    using U = typename std::remove_extent<T>::type;
    std::size_t N = std::extent<T>::value;
    return stoppable_ptr<U>(new U[N], [](U* ptr){ delete[] ptr; }); 
}

template<typename T>
class weak_stoppable_ptr
{
public:
    constexpr weak_stoppable_ptr() noexcept : ptr{nullptr}, ctrl{nullptr} {}
 
    weak_stoppable_ptr(const weak_stoppable_ptr& r) noexcept { Copy(r.ctrl, r.ptr); }

    template<typename U> 
    weak_stoppable_ptr(const weak_stoppable_ptr<U>& r) noexcept { Copy(r.ctrl, r.ptr); }

    weak_stoppable_ptr(weak_stoppable_ptr&& r) noexcept { Move(std::move(r)); }

    template<typename U> 
    weak_stoppable_ptr(weak_stoppable_ptr<U>&& r) noexcept { Move(std::move(r)); }

    template<typename U>
    weak_stoppable_ptr(const stoppable_ptr<U>& r) noexcept { Copy(r.ctrl, r.ptr); }

    ~weak_stoppable_ptr() { reset(); }

    weak_stoppable_ptr& operator=(const weak_stoppable_ptr& r) noexcept
    { return operator= <T>(r); }

    template<typename U>
    weak_stoppable_ptr& operator=(const weak_stoppable_ptr<U>& r) noexcept
    {
        weak_stoppable_ptr(r).swap(*this);
        return *this;
    }

    weak_stoppable_ptr& operator=(weak_stoppable_ptr&& r) noexcept
    { return operator= <T>(std::move(r)); }

    template<typename U>
    weak_stoppable_ptr& operator=(weak_stoppable_ptr<U>&& r) noexcept
    {
        weak_stoppable_ptr(std::move(r)).swap(*this)
        return *this;
    }

    template<typename U> 
    weak_stoppable_ptr& operator=(const stoppable_ptr<U>& r) noexcept
    {
        weak_stoppable_ptr(r).swap(*this);
        return *this;
    }

    void reset() noexcept
    {
        if (ptr)
            ptr = nullptr;

        if (!ctrl)
            return;

        if (ctrl->DecWeakRef())
            delete ctrl;

        ctrl = nullptr;
        return;
    }

    void stop() const noexcept { if (ctrl) ctrl->Stop(); }
    bool is_stopped() const noexcept { return ctrl ? ctrl->IsStopped() : false; }

    std::uint64_t use_count() const noexcept
    {
        if (!ctrl)
            return 0;

        std::lock_guard<std::mutex> lock(ctrl->mtx);
        return ctrl->strongRefCount;
    }

    bool expired() const noexcept
    {
        if (!ctrl)
            return true;

        std::lock_guard<std::mutex> lock(ctrl->mtx);
        return (ctrl->strongRefCount == 0 || ctrl->stopped);
    }

    void swap(weak_stoppable_ptr& r) noexcept
    {
        if (ctrl != r.ctrl)
            std::swap(ctrl, r.ctrl);
       
        if (ptr != r.ptr)
            std::swap(ptr, r.ptr);

        return;
    }

    stoppable_ptr<T> lock() const noexcept
    {
        if (!ctrl || ctrl->IncStrongRef() == false)
            return stoppable_ptr<T>();

        stoppable_ptr<T> out;
        out.ctrl = ctrl;
        out.ptr  = ptr;
        return out;
    }

private:
    T* ptr;
    priv_impl_stoppable_ptr::ControlBlockBase* ctrl;

    void Copy(priv_impl_stoppable_ptr::ControlBlockBase* c, T* p) noexcept
    {
        if (c && c->IncWeakRef())
        {
            ctrl = c;
            ptr  = p;
            return;
        }

        ctrl = nullptr;
        ptr  = nullptr;
        return;
    }

    template<typename U> 
    void Move(weak_stoppable_ptr<U>&& r) noexcept
    {
        if (r.ctrl)
        {
            ctrl = r.ctrl;
            ptr  = r.ptr;

            r.ctrl = nullptr;
            r.ptr  = nullptr;
            return;
        }

        ctrl = nullptr;
        ptr  = nullptr;
        return;
    }

    template<typename> friend class stoppable_ptr;
    template<typename> friend class weak_stoppable_ptr;
};

template<typename T>
template<typename U>
stoppable_ptr<T>::stoppable_ptr(const weak_stoppable_ptr<U>& r) noexcept
{
    if (r.ctrl == nullptr || r.ctrl->IncStrongRef() == false)
    {
        ctrl = nullptr;
        ptr  = nullptr;
        return;
    }

    ctrl = r.ctrl;
    ptr  = r.ptr;
    return;
}

template<typename T>
class enable_stoppable_from_this
{
public:
    using priv_impl_esft_type = T;

    stoppable_ptr<T> stoppable_from_this() { return stoppable_ptr<T>(weak_this); }
    stoppable_ptr<const T> stoppable_from_this() const { return stoppable_ptr<const T>(weak_this); }
    weak_stoppable_ptr<T> weak_stoppable_from_this() noexcept { return weak_this; }
    weak_stoppable_ptr<const T> weak_stoppable_from_this() const noexcept { return weak_this; }

protected:
    constexpr enable_stoppable_from_this() noexcept : weak_this{} {}
    enable_stoppable_from_this(const enable_stoppable_from_this<T>&) noexcept : weak_this{} {}
    ~enable_stoppable_from_this() = default;

    enable_stoppable_from_this& operator=(const enable_stoppable_from_this&) noexcept
    { return *this; }

private:
    weak_stoppable_ptr<T> weak_this;

    template<typename, typename> friend class priv_impl_stoppable_ptr::HelperSetWeakThis; 
};

} // End of namespace util.