#include "data_source.h"

#include <google/protobuf/text_format.h>

#include <saas/tools/devops/lb_dc_checker/proto/result.pb.h>
#include <saas/util/json/json.h>


namespace NFusion {
    namespace NDatacenterChecker {
        TExternalAPIDataSource::TExternalAPIDataSource(const TExternalAPIDataSourceConfig& config)
            : Config(config)
            , Request(config.Query)
        {
            Request.SetTimeout(config.Timeout.MilliSeconds());
        }

        TMap<TString, bool> TExternalAPIDataSource::ParseReplyContent(NUtil::THttpReply& reply) {
            const auto& jsonReply = NUtil::JsonFromString(reply.Content());

            TMap<TString, bool> result;
            for (const auto& item: jsonReply.GetMap()) {
                const TString& dcLabel = item.first;
                bool isFailed = !item.second["write-enabled"].GetBooleanSafe();
                result[dcLabel] = isFailed;
            }
            return result;
        }

        TCheckResult TExternalAPIDataSource::RunCheck() {
            NUtil::THttpReply reply = Request.Execute(Config.Host, Config.Port);
            if (reply.IsSuccessReply()) {
                try {
                    return TSuccessCheck(ParseReplyContent(reply));
                } catch (...) {
                    return TErrorCheck(CurrentExceptionMessage());
                }
            } else {
                return TErrorCheck(reply.ErrorMessage());
            }
        }

        TZookeeperDataSource::TZookeeperDataSource(const TZookeeperDataSourceConfig& config)
            : TZKClient(config.ZKServers, nullptr, config.ZKRequestsMaxAttempts, config.ZKLostConnectionWaitTime)
            , Config(config)
        { }

        TMap<TString, bool> TZookeeperDataSource::ParseZKData(const TString& data) {
            NSaasLB::TDcCheckerResult proto;
            Y_ENSURE(google::protobuf::TextFormat::ParseFromString(data, &proto), "Failed to parse ZK check result");

            const auto& currentTs = TInstant::Now().Seconds();
            const auto& tsDiff = TDuration::Seconds(currentTs - proto.GetUpdatedTs());
            if (tsDiff > Config.ZKMaxValidDataDelay) {
                ythrow yexception() << "ZK data is outdated, ts diff: " << tsDiff << Endl;
            }

            TMap<TString, bool> result;
            for (const auto& dc: proto.GetAvailableDcs()) {
                result[dc] = false;
            }
            for (const auto& dc: proto.GetUnavailableDcs()) {
                result[dc] = true;
            }
            return result;
        }

        TCheckResult TZookeeperDataSource::RunCheck() {
            if (!IsConnected()) {
                Reconnect();
            }

            try {
                return TSuccessCheck(ParseZKData(GetData(Config.ZKPath)));
            } catch(...) {
                return TErrorCheck(CurrentExceptionMessage());
            }
        }
    }
}
