#include "EventSink.h"

namespace {

// Turns a dispid into the event name the type library gives it, so the caller
// works in names rather than in numbers the OCX is free to renumber.
class EventSink : public IDispatch {
public:
    EventSink(ITypeInfo* events, SinkHandler handler)
        : events_(events), handler_(std::move(handler)) {
        if (events_ != nullptr) {
            events_->AddRef();
        }
    }

    virtual ~EventSink() {
        if (events_ != nullptr) {
            events_->Release();
        }
    }

    STDMETHODIMP QueryInterface(REFIID iid, void** out) override {
        if (iid == IID_IUnknown || iid == IID_IDispatch || iid == source_) {
            *out = static_cast<IDispatch*>(this);
            AddRef();
            return S_OK;
        }
        *out = nullptr;
        return E_NOINTERFACE;
    }

    STDMETHODIMP_(ULONG) AddRef() override { return ++references_; }

    STDMETHODIMP_(ULONG) Release() override {
        ULONG left = --references_;
        if (left == 0) {
            delete this;
        }
        return left;
    }

    STDMETHODIMP GetTypeInfoCount(UINT* count) override {
        *count = 0;
        return S_OK;
    }

    STDMETHODIMP GetTypeInfo(UINT, LCID, ITypeInfo**) override { return E_NOTIMPL; }

    STDMETHODIMP GetIDsOfNames(REFIID, LPOLESTR*, UINT, LCID, DISPID*) override {
        return E_NOTIMPL;
    }

    STDMETHODIMP Invoke(DISPID dispid, REFIID, LCID, WORD, DISPPARAMS* arguments, VARIANT*,
                        EXCEPINFO*, UINT*) override {
        handler_(nameOf(dispid), arguments);
        return S_OK;
    }

    void setSource(REFIID source) { source_ = source; }

private:
    std::wstring nameOf(DISPID dispid) {
        BSTR name = nullptr;
        if (events_ == nullptr ||
            FAILED(events_->GetDocumentation(dispid, &name, nullptr, nullptr, nullptr))) {
            return L"";
        }
        std::wstring found(name, SysStringLen(name));
        SysFreeString(name);
        return found;
    }

    ITypeInfo* events_;
    SinkHandler handler_;
    IID source_ = IID_NULL;
    ULONG references_ = 1;
};

}  // namespace

ITypeInfo* sourceTypeInfo(IDispatch* control, IID* source) {
    IProvideClassInfo2* provider = nullptr;
    if (FAILED(control->QueryInterface(IID_IProvideClassInfo2,
                                       reinterpret_cast<void**>(&provider)))) {
        return nullptr;
    }
    ITypeInfo* events = nullptr;
    if (SUCCEEDED(provider->GetGUID(GUIDKIND_DEFAULT_SOURCE_DISP_IID, source))) {
        ITypeInfo* coclass = nullptr;
        if (SUCCEEDED(provider->GetClassInfo(&coclass))) {
            TYPEATTR* attributes = nullptr;
            if (SUCCEEDED(coclass->GetTypeAttr(&attributes))) {
                for (UINT i = 0; i < attributes->cImplTypes && events == nullptr; i++) {
                    HREFTYPE reference = 0;
                    ITypeInfo* implemented = nullptr;
                    if (FAILED(coclass->GetRefTypeOfImplType(i, &reference)) ||
                        FAILED(coclass->GetRefTypeInfo(reference, &implemented))) {
                        continue;
                    }
                    TYPEATTR* theirs = nullptr;
                    if (SUCCEEDED(implemented->GetTypeAttr(&theirs))) {
                        if (theirs->guid == *source) {
                            events = implemented;
                            implemented->AddRef();
                        }
                        implemented->ReleaseTypeAttr(theirs);
                    }
                    implemented->Release();
                }
                coclass->ReleaseTypeAttr(attributes);
            }
            coclass->Release();
        }
    }
    provider->Release();
    return events;
}

IDispatch* createEventSink(ITypeInfo* events, REFIID source, SinkHandler handler) {
    EventSink* sink = new EventSink(events, std::move(handler));
    sink->setSource(source);
    return sink;
}