/*
 * All or portions of this file Copyright (c) Amazon.com, Inc. or its affiliates or
 * its licensors.
 *
 * For complete copyright and license terms please see the LICENSE at the root of this
 * distribution (the "License"). All use of this software is governed by the License,
 * or, if provided, by the license below or the license accompanying this file. Do not
 * remove or modify any license notices. This file is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *
 */

#include <AzCore/Jobs/JobFunction.h>

#include <UserManagement/AWSCognitoUserManagementController.h>
#include <AWSClientAuthBus.h>
#include <AWSCoreBus.h>

#include <aws/core/utils/Outcome.h>
#include <aws/core/utils/memory/stl/AWSVector.h>
#include <aws/cognito-idp/model/SignUpRequest.h>
#include <aws/cognito-idp/CognitoIdentityProviderClient.h>
#include <aws/cognito-idp/model/SignUpResult.h>
#include <aws/cognito-idp/model/ConfirmSignUpRequest.h>
#include <aws/cognito-idp/model/ConfirmSignUpResult.h>
#include <aws/cognito-idp/model/ConfirmSignUpRequest.h>
#include <aws/cognito-idp/model/ConfirmSignUpResult.h>
#include <aws/cognito-idp/model/AttributeType.h>
#include <aws/cognito-idp/model/ForgotPasswordRequest.h>
#include <aws/cognito-idp/model/ForgotPasswordResult.h>
#include <aws/cognito-idp/model/ConfirmForgotPasswordRequest.h>
#include <aws/cognito-idp/model/ConfirmForgotPasswordResult.h>
#include <aws/cognito-idp/model/SetUserMFAPreferenceRequest.h>
#include <aws/cognito-idp/model/SetUserMFAPreferenceResult.h>

namespace AWSClientAuth
{
    constexpr char COGNITO_USER_POOL[] = "/AWS/CognitoUserPool";

    AWSCognitoUserManagementController::AWSCognitoUserManagementController()
    {
        AZ::Interface<IAWSCognitoUserManagementRequests>::Register(this);
        AWSCognitoUserManagementRequestBus::Handler::BusConnect();

        m_settings = AZStd::make_unique<AWSCognitoUserManagementSetting>();
    }

    AWSCognitoUserManagementController::~AWSCognitoUserManagementController()
    {
        m_settings.reset();

        AWSCognitoUserManagementRequestBus::Handler::BusDisconnect();
        AZ::Interface<IAWSCognitoUserManagementRequests>::Unregister(this);
    }

    bool AWSCognitoUserManagementController::Initialize(const AZStd::string& settingsRegistryPath)
    {
        AZStd::unique_ptr<AZ::SettingsRegistryInterface> settingsRegistry = AZStd::make_unique<AZ::SettingsRegistryImpl>();

        if (!settingsRegistry->MergeSettingsFile(settingsRegistryPath, AZ::SettingsRegistryInterface::Format::JsonMergePatch))
        {
            AZ_Error("AWSCognitoUserManagementController", true, "Failed to merge settings file for path %s", settingsRegistryPath.c_str());
            return false;
        }

        if (!settingsRegistry->GetObject(m_settings.get(), azrtti_typeid(m_settings.get()), COGNITO_USER_POOL))
        {
            AZ_Error("AWSCognitoUserManagementController", true, "Failed to get settings object for path %s", COGNITO_USER_POOL);
            return false;
        }

        return true;
    }

    // Call Cognito user pool sign up using email. Confirmation code sent to the email set.
    // Refer https://docs.aws.amazon.com/cognito/latest/developerguide/signing-up-users-in-your-app.html
    void AWSCognitoUserManagementController::EmailSignUpAsync(const AZStd::string& username, const AZStd::string& password, const AZStd::string& email)
    {
        std::shared_ptr<Aws::CognitoIdentityProvider::CognitoIdentityProviderClient> cognitoIdentityProviderClient =
            AZ::Interface<IAWSClientAuthRequests>::Get()->GetCognitoIDPClient();
        
        AZ::JobContext* jobContext = nullptr;
        AWSCore::AWSCoreRequestBus::BroadcastResult(jobContext, &AWSCore::AWSCoreRequests::GetDefaultJobContext);

        AZ::Job* emailSignUpJob = AZ::CreateJobFunction([this, cognitoIdentityProviderClient, username, password, email]()
        {
            Aws::CognitoIdentityProvider::Model::SignUpRequest signUpRequest;
            signUpRequest.SetClientId(m_settings->m_appClientId.c_str());
            signUpRequest.SetUsername(username.c_str());
            signUpRequest.SetPassword(password.c_str());

            Aws::Vector<Aws::CognitoIdentityProvider::Model::AttributeType> attributes;
            Aws::CognitoIdentityProvider::Model::AttributeType emailAttribute;
            emailAttribute.SetName("email");
            emailAttribute.SetValue(email.c_str());
            attributes.push_back(emailAttribute);
            signUpRequest.SetUserAttributes(attributes);

            Aws::CognitoIdentityProvider::Model::SignUpOutcome signUpOutcome{ cognitoIdentityProviderClient->SignUp(signUpRequest) };
            if (signUpOutcome.IsSuccess())
            {
                Aws::CognitoIdentityProvider::Model::SignUpResult signUpResult{ signUpOutcome.GetResult() };
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnEmailSignUpSuccess, signUpResult.GetUserSub().c_str());
                
            }
            else
            {
                Aws::Client::AWSError<Aws::CognitoIdentityProvider::CognitoIdentityProviderErrors> error = signUpOutcome.GetError();
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnEmailSignUpFail, error.GetMessage().c_str());
            }
        }, true, jobContext);
        emailSignUpJob->Start();
    }

    void AWSCognitoUserManagementController::PhoneSignUpAsync(const AZStd::string& username, const AZStd::string& password, const AZStd::string& phoneNumber)
    {
        std::shared_ptr<Aws::CognitoIdentityProvider::CognitoIdentityProviderClient> cognitoIdentityProviderClient =
            AZ::Interface<IAWSClientAuthRequests>::Get()->GetCognitoIDPClient();

        AZ::JobContext* jobContext = nullptr;
        AWSCore::AWSCoreRequestBus::BroadcastResult(jobContext, &AWSCore::AWSCoreRequests::GetDefaultJobContext);

        AZ::Job* phoneSignUpJob = AZ::CreateJobFunction([this, cognitoIdentityProviderClient, username, password, phoneNumber]()
        {
            Aws::CognitoIdentityProvider::Model::SignUpRequest signUpRequest;
            signUpRequest.SetClientId(m_settings->m_appClientId.c_str());
            signUpRequest.SetUsername(username.c_str());
            signUpRequest.SetPassword(password.c_str());

            Aws::Vector<Aws::CognitoIdentityProvider::Model::AttributeType> attributes;
            Aws::CognitoIdentityProvider::Model::AttributeType emailAttribute;
            emailAttribute.SetName("phone_number");
            emailAttribute.SetValue(phoneNumber.c_str());
            attributes.push_back(emailAttribute);
            signUpRequest.SetUserAttributes(attributes);

            Aws::CognitoIdentityProvider::Model::SignUpOutcome signUpOutcome{ cognitoIdentityProviderClient->SignUp(signUpRequest) };
            if (signUpOutcome.IsSuccess())
            {
                Aws::CognitoIdentityProvider::Model::SignUpResult signUpResult{ signUpOutcome.GetResult() };
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnPhoneSignUpSuccess, signUpResult.GetUserSub().c_str());

            }
            else
            {
                Aws::Client::AWSError<Aws::CognitoIdentityProvider::CognitoIdentityProviderErrors> error = signUpOutcome.GetError();
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnPhoneSignUpFail, error.GetMessage().c_str());
            }
        }, true, jobContext);
        phoneSignUpJob->Start();
    }

    // Call Cognito user pool confirm sign up using code from email/phone.
    // Refer https://docs.aws.amazon.com/cognito/latest/developerguide/signing-up-users-in-your-app.html
    void AWSCognitoUserManagementController::ConfirmSignUpAsync(const AZStd::string& username, const AZStd::string& confirmationCode)
    {
        std::shared_ptr<Aws::CognitoIdentityProvider::CognitoIdentityProviderClient> cognitoIdentityProviderClient =
            AZ::Interface<IAWSClientAuthRequests>::Get()->GetCognitoIDPClient();
 
        AZ::JobContext* jobContext = nullptr;
        AWSCore::AWSCoreRequestBus::BroadcastResult(jobContext, &AWSCore::AWSCoreRequests::GetDefaultJobContext);

        AZ::Job* confirmSignUpJob = AZ::CreateJobFunction([this, cognitoIdentityProviderClient, username, confirmationCode]()
        {
            Aws::CognitoIdentityProvider::Model::ConfirmSignUpRequest confirmSignupRequest;
            confirmSignupRequest.SetClientId(m_settings->m_appClientId.c_str());
            confirmSignupRequest.SetUsername(username.c_str());
            confirmSignupRequest.SetConfirmationCode(confirmationCode.c_str());

            Aws::CognitoIdentityProvider::Model::ConfirmSignUpOutcome confirmSignupOutcome{ cognitoIdentityProviderClient->ConfirmSignUp(confirmSignupRequest) };
            if (confirmSignupOutcome.IsSuccess())
            {
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnConfirmSignUpSuccess);
            }
            else
            {
                Aws::Client::AWSError<Aws::CognitoIdentityProvider::CognitoIdentityProviderErrors> error = confirmSignupOutcome.GetError();
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnConfirmSignUpFail, error.GetMessage().c_str());
            }
        }, true, jobContext);
        confirmSignUpJob->Start();
    }

    void AWSCognitoUserManagementController::ForgotPasswordAsync(const AZStd::string& username)
    {
        std::shared_ptr<Aws::CognitoIdentityProvider::CognitoIdentityProviderClient> cognitoIdentityProviderClient =
            AZ::Interface<IAWSClientAuthRequests>::Get()->GetCognitoIDPClient();

        AZ::JobContext* jobContext = nullptr;
        AWSCore::AWSCoreRequestBus::BroadcastResult(jobContext, &AWSCore::AWSCoreRequests::GetDefaultJobContext);

        AZ::Job* forgotPasswordJob = AZ::CreateJobFunction([this, cognitoIdentityProviderClient, username]()
        {
            Aws::CognitoIdentityProvider::Model::ForgotPasswordRequest forgotPasswordRequest;
            forgotPasswordRequest.SetClientId(m_settings->m_appClientId.c_str());
            forgotPasswordRequest.SetUsername(username.c_str());

            Aws::CognitoIdentityProvider::Model::ForgotPasswordOutcome forgotPasswordOutcome{ cognitoIdentityProviderClient->ForgotPassword(forgotPasswordRequest) };
            if (forgotPasswordOutcome.IsSuccess())
            {
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnForgotPasswordSuccess);
            }
            else
            {
                Aws::Client::AWSError<Aws::CognitoIdentityProvider::CognitoIdentityProviderErrors> error = forgotPasswordOutcome.GetError();
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnForgotPasswordFail, error.GetMessage().c_str());
            }
        }, true, jobContext);
        forgotPasswordJob->Start();
    }

    void AWSCognitoUserManagementController::ConfirmForgotPasswordAsync(const AZStd::string& username, const AZStd::string& confirmationCode, const AZStd::string& newPassword)
    {
        std::shared_ptr<Aws::CognitoIdentityProvider::CognitoIdentityProviderClient> cognitoIdentityProviderClient =
            AZ::Interface<IAWSClientAuthRequests>::Get()->GetCognitoIDPClient();

        AZ::JobContext* jobContext = nullptr;
        AWSCore::AWSCoreRequestBus::BroadcastResult(jobContext, &AWSCore::AWSCoreRequests::GetDefaultJobContext);

        AZ::Job* confirmForgotPasswordJob = AZ::CreateJobFunction([this, cognitoIdentityProviderClient, username, confirmationCode, newPassword]()
        {
            Aws::CognitoIdentityProvider::Model::ConfirmForgotPasswordRequest confirmForgotPasswordRequest;
            confirmForgotPasswordRequest.SetClientId(m_settings->m_appClientId.c_str());
            confirmForgotPasswordRequest.SetUsername(username.c_str());
            confirmForgotPasswordRequest.SetConfirmationCode(confirmationCode.c_str());
            confirmForgotPasswordRequest.SetPassword(newPassword.c_str());

            Aws::CognitoIdentityProvider::Model::ConfirmForgotPasswordOutcome confirmForgotPasswordOutcome{ cognitoIdentityProviderClient->ConfirmForgotPassword(confirmForgotPasswordRequest) };
            if (confirmForgotPasswordOutcome.IsSuccess())
            {
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnConfirmForgotPasswordSuccess);
            }
            else
            {
                Aws::Client::AWSError<Aws::CognitoIdentityProvider::CognitoIdentityProviderErrors> error = confirmForgotPasswordOutcome.GetError();
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnConfirmForgotPasswordFail, error.GetMessage().c_str());
            }
        }, true, jobContext);
        confirmForgotPasswordJob->Start();
    }

    void AWSCognitoUserManagementController::EnableMFAAsync(const AZStd::string& accessToken)
    {
        std::shared_ptr<Aws::CognitoIdentityProvider::CognitoIdentityProviderClient> cognitoIdentityProviderClient =
            AZ::Interface<IAWSClientAuthRequests>::Get()->GetCognitoIDPClient();

        AZ::JobContext* jobContext = nullptr;
        AWSCore::AWSCoreRequestBus::BroadcastResult(jobContext, &AWSCore::AWSCoreRequests::GetDefaultJobContext);

        AZ::Job* enableMFAJob = AZ::CreateJobFunction([this, cognitoIdentityProviderClient, accessToken]()
        {
            Aws::CognitoIdentityProvider::Model::SetUserMFAPreferenceRequest confirmForgotPasswordRequest;
            Aws::CognitoIdentityProvider::Model::SMSMfaSettingsType settings;
            settings.SetEnabled(true);
            settings.SetPreferredMfa(true);
            confirmForgotPasswordRequest.SetSMSMfaSettings(settings);
            confirmForgotPasswordRequest.SetAccessToken(accessToken.c_str());

            Aws::CognitoIdentityProvider::Model::SetUserMFAPreferenceOutcome setUserMFAPreferenceOutcome{ cognitoIdentityProviderClient->SetUserMFAPreference(confirmForgotPasswordRequest) };
            if (setUserMFAPreferenceOutcome.IsSuccess())
            {
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnEnableMFASuccess);
            }
            else
            {
                Aws::Client::AWSError<Aws::CognitoIdentityProvider::CognitoIdentityProviderErrors> error = setUserMFAPreferenceOutcome.GetError();
                AWSCognitoUserManagementNotificationBus::Broadcast(&AWSCognitoUserManagementNotifications::OnEnableMFAFail, error.GetMessage().c_str());
            }
        }, true, jobContext);
        enableMFAJob->Start();
    }
} // namespace AWSClientAuth