#include "dc_checker.h"


namespace NFusion {
    TDatacenterChecker::TDatacenterChecker(const TDatacenterCheckerConfig& config)
        : Config(config)
    {
        if (config.DataSourceType == EDataSourceType::ExternalAPI) {
            DataSource = MakeHolder<NDatacenterChecker::TExternalAPIDataSource>(
                std::get<TExternalAPIDataSourceConfig>(config.DataSource)
            );
        } else if (config.DataSourceType == EDataSourceType::Zookeeper) {
            DataSource = MakeHolder<NDatacenterChecker::TZookeeperDataSource>(
                std::get<TZookeeperDataSourceConfig>(config.DataSource)
            );
        } else {
            ythrow yexception() << "Unexpected data source type";
        }

        RegisterGlobalMessageProcessor(this);
    }

    TDatacenterChecker::~TDatacenterChecker() {
        UnregisterGlobalMessageProcessor(this);
    }

    void TDatacenterChecker::Start() {
        CHECK_WITH_LOG(!Checker);

        TThread::TParams checkerThreadParams(&CheckerProc, this);
        checkerThreadParams.SetName("DfDcCheck");

        Checker = MakeHolder<TThread>(checkerThreadParams);
        Checker->Start();

        INFO_LOG << "Started datacenter checker thread for streams: " << Config.StreamType << Endl;
    }

    void TDatacenterChecker::Stop() {
        NeedStop.Signal();
        if (Checker) {
            Checker->Join();
            Checker.Destroy();
            INFO_LOG << "Stopped datacenter checker thread for streams: " << Config.StreamType << Endl;
        }

        TGuard <TMutex> guard(Lock);
        Subscribers.clear();
        Statuses.clear();
        NeedStop.Reset();
    }

    void TDatacenterChecker::Subscribe(ISubscriber *subscriber) {
        TGuard <TMutex> guard(Lock);
        Subscribers.push_back(subscriber);
    }

    TString TDatacenterChecker::GetCorrectDatacenter(const TVector<TString> datacenters) const {
        TGuard <TMutex> guard(Lock);
        for (const auto& dc: datacenters) {
            const auto& it = Statuses.find(dc);
            if (it == Statuses.end() || !it->second.IsFailed) {
                return dc;
            }
        }
        return datacenters.front();
    }

    NJson::TJsonValue TDatacenterChecker::GetStatus() const {
        TGuard <TMutex> guard(Lock);
        NJson::TJsonValue status;
        for (const auto& dc: Statuses) {
            NJson::TJsonValue dcStatus;
            dcStatus.InsertValue("failed", dc.second.IsFailed);
            dcStatus.InsertValue("blockedUntil", dc.second.BlockedUntil.ToStringLocal());
            status.InsertValue(dc.first, dcStatus);
        }
        return status;
    }

    TString TDatacenterChecker::Name() const {
        return "DatacenterChecker";
    }

    bool TDatacenterChecker::Process(IMessage *message_) {
        if (auto message = message_->As<TMessageGetDocfetcherStatus>()) {
            message->Status["DatacenterChecker"] = GetStatus();
            return true;
        }
        if (auto message = message_->As<TMessageSetDocfetcherDatacentersStatus>()) {
            TGuard <TMutex> guard(Lock);
            for (const auto& dcLabel: message->Datacenters) {
                SetStatus(dcLabel, Statuses[dcLabel], message->Failed, message->BlockTime);
            }
            return true;
        }
        return false;
    }

    void TDatacenterChecker::SendNotify(const TString& dc, bool isFailed) const {
        for (auto& subscriber : Subscribers) {
            subscriber->DatacenterChangedStatus(dc, isFailed);
        }
    }

    void TDatacenterChecker::SetStatus(
        const TString& dc,
        NDatacenterChecker::TDatacenterStatus& status,
        bool isFailed,
        const TDuration& blockDuration
    ) {
        bool statusChanged = status.IsFailed != isFailed;
        status.IsFailed = isFailed;
        status.BlockedUntil = TInstant::Now() + blockDuration;
        if (statusChanged) {
            INFO_LOG << "Changed status of the datacenter " << dc << " : " << (isFailed ? "FAILED" : "OK") << Endl;
            SendNotify(dc, isFailed);
        }
    }

    void TDatacenterChecker::UpdateStatus(
        const TString& dc,
        NDatacenterChecker::TDatacenterStatus& status,
        bool isFailed
    ) {
        if (TInstant::Now() < status.BlockedUntil) {
            return;
        }
        SetStatus(dc, status, isFailed, isFailed ? Config.BlockDuration : TDuration::Seconds(0));
    }

    void TDatacenterChecker::TryClearStatuses() {
        TGuard<TMutex> guard(Lock);
        for (auto& dc : Statuses) {
            UpdateStatus(dc.first, dc.second, false);
        }
    }

    void TDatacenterChecker::ApplyCheckResult(const NDatacenterChecker::TSuccessCheck& successCheck) {
        TGuard<TMutex> guard(Lock);
        for (const auto& [dc, failed] : successCheck.DcToFailed) {
            UpdateStatus(dc, Statuses[dc], failed);
        }
    }

    void* TDatacenterChecker::CheckerProc(void* object) {
        ThreadDisableBalloc();
        auto this_ = reinterpret_cast<TDatacenterChecker*>(object);
        const auto& config = this_->Config;
        auto& needStop = this_->NeedStop;

        ui32 consecutiveErrors = 0;
        do {
            const NDatacenterChecker::TCheckResult& result = this_->DataSource->RunCheck();
            TString error;

            if (std::holds_alternative<NDatacenterChecker::TSuccessCheck>(result)) {
                const auto& successCheck = std::get<NDatacenterChecker::TSuccessCheck>(result);

                try {
                    this_->ApplyCheckResult(successCheck);
                } catch (...) {
                    error = "Unable to apply the result of checking data centers: " + CurrentExceptionMessage();
                }
            } else {
                const auto& errorCheck = std::get<NDatacenterChecker::TErrorCheck>(result);
                error = "Unable to get availability statuses of data centers. Last error: " + errorCheck.ErrorMessage;
            }

            if (error) {
                ++consecutiveErrors;
                if (consecutiveErrors > config.MaxConsecutiveErrors) {
                    consecutiveErrors = 0;
                    ERROR_LOG << error << Endl;
                    this_->TryClearStatuses();
                }
            }
        } while (!needStop.WaitT(config.CheckInterval));
        return nullptr;
    }
}
