/*
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#pragma once

#include <AzCore/Component/TickBus.h>

#include <Framework/AWSApiClientJob.h>
#include <Framework/AWSApiRequestJobConfig.h>


namespace Aws
{
    namespace Client
    {
        class AsyncCallerContext;
    }
}

namespace AWSCore
{

    /// A macro that simplifies the declaration of AwsApiRequestTraits for APIs that have a result.
    ///
    /// \param SERVICE - the name of an Aws service. e.g. Lambda.
    ///
    /// \param REQUEST - the name of an Aws request supported by that 
    /// service. e.g. Invoke.
#define AWS_API_REQUEST_TRAITS(SERVICE, REQUEST) \
        AWSCore::AwsApiRequestTraits< \
            Aws::SERVICE::SERVICE##Client, \
            Aws::SERVICE::Model::REQUEST##Request, \
            Aws::SERVICE::Model::REQUEST##Outcome, \
            Aws::SERVICE::Model::REQUEST##OutcomeCallable, \
            Aws::SERVICE::REQUEST##ResponseReceivedHandler, \
            Aws::SERVICE::Model::REQUEST##Result, \
            Aws::Client::AWSError<Aws::SERVICE::SERVICE##Errors>, \
            &Aws::SERVICE::SERVICE##Client::REQUEST, \
            &Aws::SERVICE::SERVICE##Client::REQUEST##Callable, \
            &Aws::SERVICE::SERVICE##Client::REQUEST##Async \
        >

    /// A macro that simplifies the declaration of AwsApiRequestTraits for APIs that have no result.
    ///
    /// \param SERVICE - the name of an Aws service. e.g. Lambda.
    ///
    /// \param REQUEST - the name of an Aws request supported by that 
    /// service. e.g. Invoke.
#define AWS_API_REQUEST_TRAITS_NO_RESULT(SERVICE, REQUEST) \
        AWSCore::AwsApiRequestTraits< \
            Aws::SERVICE::SERVICE##Client, \
            Aws::SERVICE::Model::REQUEST##Request, \
            Aws::SERVICE::Model::REQUEST##Outcome, \
            Aws::SERVICE::Model::REQUEST##OutcomeCallable, \
            Aws::SERVICE::REQUEST##ResponseReceivedHandler, \
            Aws::NoResult, \
            Aws::Client::AWSError<Aws::SERVICE::SERVICE##Errors>, \
            &Aws::SERVICE::SERVICE##Client::REQUEST, \
            &Aws::SERVICE::SERVICE##Client::REQUEST##Callable, \
            &Aws::SERVICE::SERVICE##Client::REQUEST##Async \
        >

    /// Macro used below to define the traits type. Parameters
    /// correspond to those generated by the macro above.
#define AWS_API_REQUEST_TRAITS_TEMPLATE_DEFINITION_HELPER \
        template< \
            class _ClientType, \
            class _RequestType, \
            class _OutcomeType, \
            class _OutcomeCallableType, \
            class _AsyncHandlerType, \
            class _ResultType, \
            class _ErrorType, \
            RequestFunctionType<_ClientType, _OutcomeType, _RequestType> _Function, \
            RequestCallableFunctionType<_ClientType, _OutcomeCallableType, _RequestType> _CallableFunction, \
            RequestAsyncFunctionType<_ClientType, _RequestType, _AsyncHandlerType> _AsyncFunction \
        >

    /// Macro used below to define the traits type. Parameters
    /// correspond to those generated by the macro above.
#define AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER \
        AwsApiRequestTraits< \
            _ClientType, \
            _RequestType, \
            _OutcomeType, \
            _OutcomeCallableType, \
            _AsyncHandlerType, \
            _ResultType, \
            _ErrorType, \
            _Function, \
            _CallableFunction, \
            _AsyncFunction \
        >

    // const pointer to const sync execution member function of client type
    template<class ClientType, class OutcomeType, class RequestType>
    using RequestFunctionType = OutcomeType(ClientType::* const)(const RequestType&) const;

    // const pointer to const callable execution member function of client type
    template<class ClientType, class OutcomeCallableType, class RequestType>
    using RequestCallableFunctionType = OutcomeCallableType(ClientType::* const)(const RequestType&) const;

    // const pointer to const async execution member function of client type
    template<class ClientType, class RequestType, class AsyncHandlerType>
    using RequestAsyncFunctionType = void (ClientType::* const)(const RequestType&, const AsyncHandlerType&, const std::shared_ptr<const Aws::Client::AsyncCallerContext>&) const;

    /// Encapsulates everything we need to know to make an AWS service
    /// request of a particular type. Use the AWS_API_REQUEST_TRAITS macro 
    /// instead of using this type directly.
    AWS_API_REQUEST_TRAITS_TEMPLATE_DEFINITION_HELPER
        struct AwsApiRequestTraits
    {

        using ClientType = _ClientType;
        using RequestType = _RequestType;
        using OutcomeType = _OutcomeType;
        using CallableOutcomeType = _OutcomeCallableType;
        using AsyncHandlerType = _AsyncHandlerType;
        using ResultType = _ResultType;
        using ErrorType = _ErrorType;
        using FunctionType = RequestFunctionType<_ClientType, _OutcomeType, _RequestType>;
        using CallableFunctionType = RequestCallableFunctionType<_ClientType, _OutcomeCallableType, _RequestType>;
        using AsyncFunctionType = RequestAsyncFunctionType<_ClientType, _RequestType, _AsyncHandlerType>;

        static FunctionType Function;
        static CallableFunctionType CallableFunction;
        static AsyncFunctionType AsyncFunction;

    };

    /// Define storage for AwsApiRequestTraits::Function
    AWS_API_REQUEST_TRAITS_TEMPLATE_DEFINITION_HELPER
        typename AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER::FunctionType AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER::Function = _Function;

    /// Define storage for AwsApiRequestTraits::CallableFunction
    AWS_API_REQUEST_TRAITS_TEMPLATE_DEFINITION_HELPER
        typename AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER::CallableFunctionType AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER::CallableFunction = _CallableFunction;

    /// Define storage for AwsApiRequestTraits::AsyncFunction
    AWS_API_REQUEST_TRAITS_TEMPLATE_DEFINITION_HELPER
        typename AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER::AsyncFunctionType AWS_API_REQUEST_TRAITS_TEMPLATE_INSTANCE_HELPER::AsyncFunction = _AsyncFunction;

    /// Macro that simplifies the declaration of an AwsRequestJob that has a result.
#define AWS_API_REQUEST_JOB(SERVICE, REQUEST) AWSCore::AwsApiRequestJob<AWS_API_REQUEST_TRAITS(SERVICE, REQUEST)>

/// Macro that simplifies the declaration of an AwsRequestJob that has no result.
#define AWS_API_REQUEST_JOB_NO_RESULT(SERVICE, REQUEST) AWSCore::AwsApiRequestJob<AWS_API_REQUEST_TRAITS_NO_RESULT(SERVICE, REQUEST)>

/// An Az::Job that that executes a specific AWS request.
    template<class RequestTraits>
    class AwsApiRequestJob
        : public AwsApiClientJob<typename RequestTraits::ClientType>
    {

    public:

        // To use a different allocator, extend this class and use this macro.
        AZ_CLASS_ALLOCATOR(AwsApiRequestJob, AZ::SystemAllocator);

        using AwsApiRequestJobType = AwsApiRequestJob<RequestTraits>;
        using AwsApiClientJobType = AwsApiClientJob<typename RequestTraits::ClientType>;

        using ClientType = typename RequestTraits::ClientType;
        using RequestType = typename RequestTraits::RequestType;
        using OutcomeType = typename RequestTraits::OutcomeType;
        using ResultType = typename RequestTraits::ResultType;
        using ErrorType = typename RequestTraits::ErrorType;

        using IConfig = IAwsApiRequestJobConfig<RequestTraits>;
        using Config = AwsApiRequestJobConfig<RequestTraits>;

        using OnSuccessFunction = AZStd::function<void(AwsApiRequestJob* job)>;
        using OnFailureFunction = AZStd::function<void(AwsApiRequestJob* job)>;

        class Function;

        template<class Allocator = AZ::SystemAllocator>
        static AwsApiRequestJob* Create(OnSuccessFunction onSuccess, OnFailureFunction onFailure = OnFailureFunction{}, IConfig* config = GetDefaultConfig());

        static Config* GetDefaultConfig()
        {
            static AwsApiJobConfigHolder<Config> s_configHolder{};
            return s_configHolder.GetConfig(AwsApiClientJobType::GetDefaultConfig());
        }

        AwsApiRequestJob(bool isAutoDelete, IConfig* config = GetDefaultConfig())
            : AwsApiClientJobType(isAutoDelete, config)
        {
        }

        /// Override AZ:Job defined method to reset request state when
        /// the job object is reused.
        void Reset(bool isClearDependent) override
        {
            request = RequestType{};
            result = ResultType{};
            error = ErrorType{};
            m_wasSuccess = false;
            AwsApiClientJobType::Reset(isClearDependent);
        }

        /// Determines if the request was successful.
        bool WasSuccess()
        {
            return m_wasSuccess;
        }

        RequestType request;
        ResultType result;
        ErrorType error;

    protected:
        void Process() override
        {

            // Currently the AWS SDK always executes requests synchronously. Even
            // when the Async and Callable versions of the API are used, the 
            // actual AWS call is made synchronously, blocking whatever thread it
            // is executing on. Eventually the SDK may use true non-blocking
            // async I/O for requests. When that feature is available, we can 
            // use the AZ::Job defined IncrementDependentCount method, start the 
            // async i/o, and call WaitForChildren. When the i/o completes, we 
            // would call DecrementDependentCount, which would cause WaitFor-
            // Children to return. We would then call OnOutcome.
            //
            // We want to call OnOutcome when the job is in a processing state 
            // so that it can create and wait for child jobs, or do other 
            // things that the job system requires it to be in the processing 
            // state. That means that in the async case it has to be called 
            // after the call to WaitForChildren returns (otherwise the job's
            // state is suspended).

            bool ok = PrepareRequest();
            if (ok)
            {
                // Get real pointer from shared pointer, then use with member
                // function pointer to call the function.
                OutcomeType outcome = (AwsApiClientJobType::m_client.get()->*RequestTraits::Function)(request);

                if (outcome.IsSuccess())
                {
                    result = std::move(outcome.GetResultWithOwnership());
                    m_wasSuccess = true;
                    OnSuccess();
                }
                else
                {
                    error = outcome.GetError();
                    m_wasSuccess = false;
                    OnFailure();
                }
            }

        }

        /// Called by Process to prepare the request. By default no changes
        /// are made to the request object. Override to defer the preparation
        /// of request data until your running on the job's worker thread,
        /// instead of setting the request data before calling Start.
        ///
        /// \return true if the request should be made.
        virtual bool PrepareRequest()
        {
            return true;
        }

        /// Called when request has completed successfully.
        virtual void OnSuccess()
        {
        }

        /// Called when the request fails.
        virtual void OnFailure()
        {
        }

        /// Called when request can't process and still requires cleanup (Specifically for the derived class AwsApiRequestJob<RequestTraits>::Function which does not use auto delete)
        virtual void DoCleanup()
        {
        }

        bool m_wasSuccess{ false };
    };

    /// A specialization of AwsApiRequestJob that lets you provide functions
    /// that are called on success or failure of the request.
    template<class RequestTraits>
    class AwsApiRequestJob<RequestTraits>::Function 
        : public AwsApiRequestJob
    {
    public:
        // To use a different allocator, extend this class and use this macro.
        AZ_CLASS_ALLOCATOR(Function, AZ::SystemAllocator);

        Function(OnSuccessFunction onSuccess, OnFailureFunction onFailure = OnFailureFunction{}, IConfig* config = GetDefaultConfig())
            : AwsApiRequestJobType(
                  false, config) // No auto delete - we need to perform our callbacks on the main thread so we queue them through tickbus
            , m_onSuccess{ onSuccess }
            , m_onFailure{ onFailure }
        {
        }

    private:
        void OnSuccess() override
        {
            AZStd::function<void()> callbackHandler = [this]()
            {
                if (m_onSuccess)
                {
                    m_onSuccess(this);
                }
                delete this;
            };
            AZ::TickBus::QueueFunction(callbackHandler);
        }

        void OnFailure() override
        {
            AZStd::function<void()> callbackHandler = [this]()
            {
                if (m_onFailure)
                {
                    m_onFailure(this);
                }
                delete this;
            };
            AZ::TickBus::QueueFunction(callbackHandler);
        }

        // Code doesn't use auto delete - this allows code to make sure things get cleaned up in cases where success or failure can't be
        // called.
        void DoCleanup() override
        {
            AZStd::function<void()> callbackHandler = [this]()
            {
                delete this;
            };
            AZ::TickBus::QueueFunction(callbackHandler);
        }

        OnSuccessFunction m_onSuccess;
        OnFailureFunction m_onFailure;
    };

    template<class RequestTraits>
    template<class Allocator>
    AwsApiRequestJob<RequestTraits>* AwsApiRequestJob<RequestTraits>::Create(
        OnSuccessFunction onSuccess, OnFailureFunction onFailure, IConfig* config)
    {
        return azcreate(Function, (onSuccess, onFailure, config), Allocator);
    }
} // namespace AWSCore
