#ifndef SLANG_COM_OBJECT_H
#define SLANG_COM_OBJECT_H

#include "slang-basic.h"

#include <atomic>

namespace Slang
{

/// A base class for COM interfaces that require atomic ref counting
/// and are *NOT* derived from RefObject
class ComBaseObject
{
public:
    /// If assigned the the ref count is *NOT* copied
    ComBaseObject& operator=(const ComBaseObject&) { return *this; }

    /// Copy Ctor, does not copy ref count
    ComBaseObject(const ComBaseObject&)
        : m_refCount(0)
    {
    }

    /// Default Ctor sets with no refs
    ComBaseObject()
        : m_refCount(0)
    {
    }

    /// Dtor needs to be virtual to avoid needing to
    /// Implement release for all derived types.
    virtual ~ComBaseObject() {}

protected:
    inline uint32_t _releaseImpl();

    std::atomic<uint32_t> m_refCount;
};

// ------------------------------------------------------------------
inline uint32_t ComBaseObject::_releaseImpl()
{
    // Check there is a ref count to avoid underflow
    SLANG_ASSERT(m_refCount != 0);
    const uint32_t count = --m_refCount;
    if (count == 0)
    {
        delete this;
    }
    return count;
}

#define SLANG_COM_BASE_IUNKNOWN_QUERY_INTERFACE                                                    \
    SLANG_NO_THROW SlangResult SLANG_MCALL queryInterface(SlangUUID const& uuid, void** outObject) \
        SLANG_OVERRIDE                                                                             \
    {                                                                                              \
        void* intf = getInterface(uuid);                                                           \
        if (intf)                                                                                  \
        {                                                                                          \
            ++m_refCount;                                                                          \
            *outObject = intf;                                                                     \
            return SLANG_OK;                                                                       \
        }                                                                                          \
        return SLANG_E_NO_INTERFACE;                                                               \
    }
#define SLANG_COM_BASE_IUNKNOWN_ADD_REF                         \
    SLANG_NO_THROW uint32_t SLANG_MCALL addRef() SLANG_OVERRIDE \
    {                                                           \
        return ++m_refCount;                                    \
    }
#define SLANG_COM_BASE_IUNKNOWN_RELEASE                          \
    SLANG_NO_THROW uint32_t SLANG_MCALL release() SLANG_OVERRIDE \
    {                                                            \
        return _releaseImpl();                                   \
    }
#define SLANG_COM_BASE_IUNKNOWN_ALL         \
    SLANG_COM_BASE_IUNKNOWN_QUERY_INTERFACE \
    SLANG_COM_BASE_IUNKNOWN_ADD_REF         \
    SLANG_COM_BASE_IUNKNOWN_RELEASE


/// COM object that derives from RefObject
class ComObject : public RefObject
{
protected:
    std::atomic<uint32_t> comRefCount;

public:
    ComObject()
        : comRefCount(0)
    {
    }
    ComObject(const ComObject& rhs)
        : RefObject(rhs), comRefCount(0)
    {
    }

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

    virtual void comFree() {}

    uint32_t addRefImpl()
    {
        auto oldRefCount = comRefCount++;
        if (oldRefCount == 0)
            addReference();
        return oldRefCount + 1;
    }

    uint32_t releaseImpl()
    {
        auto oldRefCount = comRefCount--;
        if (oldRefCount == 1)
        {
            comFree();
            releaseReference();
        }
        return oldRefCount - 1;
    }
};

#define SLANG_COM_OBJECT_IUNKNOWN_QUERY_INTERFACE                                                  \
    SLANG_NO_THROW SlangResult SLANG_MCALL queryInterface(SlangUUID const& uuid, void** outObject) \
        SLANG_OVERRIDE                                                                             \
    {                                                                                              \
        void* intf = getInterface(uuid);                                                           \
        if (intf)                                                                                  \
        {                                                                                          \
            addRef();                                                                              \
            *outObject = intf;                                                                     \
            return SLANG_OK;                                                                       \
        }                                                                                          \
        return SLANG_E_NO_INTERFACE;                                                               \
    }
#define SLANG_COM_OBJECT_IUNKNOWN_ADD_REF                       \
    SLANG_NO_THROW uint32_t SLANG_MCALL addRef() SLANG_OVERRIDE \
    {                                                           \
        return addRefImpl();                                    \
    }
#define SLANG_COM_OBJECT_IUNKNOWN_RELEASE                        \
    SLANG_NO_THROW uint32_t SLANG_MCALL release() SLANG_OVERRIDE \
    {                                                            \
        return releaseImpl();                                    \
    }
#define SLANG_COM_OBJECT_IUNKNOWN_ALL         \
    SLANG_COM_OBJECT_IUNKNOWN_QUERY_INTERFACE \
    SLANG_COM_OBJECT_IUNKNOWN_ADD_REF         \
    SLANG_COM_OBJECT_IUNKNOWN_RELEASE

} // namespace Slang

#endif