#include "pch.h"
#include <ctxtcall.h>
#include <winrt/Windows.System.h>

using namespace std::literals;
using namespace winrt;
using namespace Windows::Foundation;
using namespace Windows::System;

namespace
{
    IAsyncAction Action()
    {
        co_return;
    }

    IAsyncActionWithProgress<int> ActionProgress()
    {
        co_await 500ms;
        auto progress = co_await get_progress_token();
        progress(123);
        co_return;
    }

    IAsyncOperation<int> Operation()
    {
        co_return 123;
    }

    IAsyncOperationWithProgress<int, int> OperationProgress()
    {
        co_await 500ms;
        auto progress = co_await get_progress_token();
        progress(123);
        co_return 123;
    }

// Exclude on mingw-w64 to suppress `-Wunused-function`
#if !defined(__MINGW32__)
    bool is_mta()
    {
        APTTYPE type;
        APTTYPEQUALIFIER qualifier;
        check_hresult(CoGetApartmentType(&type, &qualifier));
        return type == APTTYPE_MTA;
    }
#endif
}

TEST_CASE("disconnected,handler,1")
{
    event<EventHandler<int>> source;

    source.add([](auto...)
        {
            throw hresult_error(RPC_E_DISCONNECTED);
        });

    auto token = source.add([](auto...)
        {
            throw hresult_error(E_INVALIDARG);
        });

    // Should have two delegates
    REQUIRE(source);

    // Should lose the disconnected delegate
    source(nullptr, 123);
    REQUIRE(source);

    // Fire the remaining delegate
    source(nullptr, 123);
    REQUIRE(source);

    // Remove the final delegate
    source.remove(token);

    // No more delegates
    REQUIRE(!source);

    source(nullptr, 123);
}

TEST_CASE("disconnected,handler,2")
{
    auto async = Action();

    async.Completed([](auto&&...)
        {
            throw hresult_error(RPC_E_DISCONNECTED);
        });
}

TEST_CASE("disconnected,handler,3")
{
    auto async = ActionProgress();
    handle signal{ CreateEventW(nullptr, true, false, nullptr) };

    async.Progress([](auto&&...)
        {
            throw hresult_error(RPC_E_DISCONNECTED);
        });

    async.Completed([&](auto&&...)
        {
            SetEvent(signal.get());
            throw hresult_error(RPC_E_DISCONNECTED);
        });

    WaitForSingleObject(signal.get(), INFINITE);
    // Give some time for to_hresult() to complete.
    Sleep(500);
}

TEST_CASE("disconnected,handler,4")
{
    auto async = Operation();

    async.Completed([](auto&&...)
        {
            throw hresult_error(RPC_E_DISCONNECTED);
        });
}

TEST_CASE("disconnected,handler,5")
{
    auto async = OperationProgress();
    handle signal{ CreateEventW(nullptr, true, false, nullptr) };

    async.Progress([](auto&&...)
        {
            throw hresult_error(RPC_E_DISCONNECTED);
        });

    async.Completed([&](auto&&...)
        {
            SetEvent(signal.get());
            throw hresult_error(RPC_E_DISCONNECTED);
        });

    WaitForSingleObject(signal.get(), INFINITE);
    // Give some time for to_hresult() to complete.
    Sleep(500);
}

// Custom action to simulate an out-of-process server that crashes before it can complete.
struct non_agile_abandoned_action : implements<non_agile_abandoned_action, IAsyncAction, IAsyncInfo, non_agile>
{
    non_agile_abandoned_action(delegate<> disconnect) : m_disconnect(disconnect) {}

    static fire_and_forget final_release(std::unique_ptr<non_agile_abandoned_action> self)
    {
        // The C++/WinRT m_handler is agile but not context-aware,
        // so we need to make sure to release it from the context it
        // was created from, which for this particular test is the MTA.
        co_await resume_background();
        // Now we can destruct.
    }

    void Completed(AsyncActionCompletedHandler const& handler) {
        m_handler = handler;
        m_disconnect();
    }
    auto Completed() { return m_handler; }
    void GetResults() {}

    auto Id() { return 0U; }
    auto Status() { return AsyncStatus::Completed; }
    auto ErrorCode() { return hresult(0); }
    void Cancel() {}
    void Close() {}

    AsyncActionCompletedHandler m_handler;
    delegate<> m_disconnect;
};

// Not yet buildable on mingw-w64.
// Missing CLSID_ContextSwitcher, IID_ICallbackWithNoReentrancyToApplicationSTA
// and __uuidof(IContextCallback). Also, the lambda needs to have __stdcall
// specified on it but there is a Clang crash bug blocking this:
// https://github.com/llvm/llvm-project/issues/58366
#if !defined(__MINGW32__)
namespace
{
    template<typename TLambda>
    void InvokeInContext(IContextCallback* context, TLambda&& lambda)
    {
        ComCallData data;
        data.pUserDefined = &lambda;
        check_hresult(context->ContextCallback([](ComCallData* data) -> HRESULT
            {
                auto& lambda = *reinterpret_cast<TLambda*>(data->pUserDefined);
                lambda();
                return S_OK;
            }, &data, IID_ICallbackWithNoReentrancyToApplicationSTA, 5, nullptr));
    }

    fire_and_forget disconnect_on_signal(com_ptr<IContextCallback> context, void* signal)
    {
        co_await resume_on_signal(signal);
        InvokeInContext(context.get(), []()
            {
                // This disconnects the IAsyncAction, simulating a server crash.
                CoDisconnectContext(INFINITE);
            });
    }
}

struct holds_hresult : public Catch::MatcherBase<hresult_error>
{
    holds_hresult(hresult value) : expected(value) {}

    hresult expected;

    bool match(hresult_error const& e) const override
    {
        return e.code() == expected;
    }

    virtual std::string describe() const override
    {
        return "is code " + std::to_string(expected.value);
    }
};

TEST_CASE("disconnected,action")
{
    auto private_context = create_instance<IContextCallback>(CLSID_ContextSwitcher);
    handle signal{ CreateEventW(nullptr, true, false, nullptr) };
    disconnect_on_signal(private_context, signal.get());

    agile_ref<IAsyncAction> action;
    InvokeInContext(private_context.get(), [&]()
        {
            action = make<non_agile_abandoned_action>([&]{ SetEvent(signal.get()); });
        });

    auto result = [](IAsyncAction action) -> IAsyncAction
        {
            co_await action;
        }(action.get());

    REQUIRE_THROWS_MATCHES(result.get(), hresult_error, holds_hresult(RPC_E_DISCONNECTED));
}

#if defined(__clang__) && (defined(_M_IX86) || defined(__i386__))
// FIXME: Test is known to crash with exit code 0xc000070a on x86 when built with Clang.
TEST_CASE("disconnected,double", "[.clang-crash]")
#else
TEST_CASE("disconnected,double")
#endif
{
    // The double-disconnect case, where the IAsyncAction disconnects,
    // and tries to return to the original context, but it too has disconnected!
    auto test = []() -> IAsyncAction
    {
        auto private_context = create_instance<IContextCallback>(CLSID_ContextSwitcher);
        handle signal{ CreateEventW(nullptr, true, false, nullptr) };
        disconnect_on_signal(private_context, signal.get());

        // Create an STA thread that we will destroy while awaiting.
        auto controller = DispatcherQueueController::CreateOnDedicatedThread();

        agile_ref<IAsyncAction> action;
        InvokeInContext(private_context.get(), [&]()
            {
                action = make<non_agile_abandoned_action>([&]() -> fire_and_forget
                    {
                        // Get off the DispatcherQueue thread.
                        co_await resume_background();
                        // Destroy the DispatcherQueue, so the co_await has nowhere to return to.
                        co_await controller.ShutdownQueueAsync();
                        // Now set the event to force the action to disconnect.
                        SetEvent(signal.get());
                    });
            });

        // Go to our STA thread.
        co_await resume_foreground(controller.DispatcherQueue());

        HRESULT hr = S_OK;
        try
        {
            co_await action.get();
        }
        catch (...)
        {
            hr = to_hresult();
            REQUIRE(is_mta());
        }
        REQUIRE(FAILED(hr));

    }();

    test.get();
}
#endif
