#include "module.h"

#include <balancer/kernel/balancer/average.h>
#include <balancer/kernel/custom_io/null.h>
#include <balancer/kernel/fs/kv_file_consumer.h>
#include <balancer/kernel/fs/shared_file_exists_checker.h>
#include <balancer/kernel/fs/shared_file_rereader.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/helpers/errors.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/http/parser/response_builder.h>
#include <balancer/kernel/http/parser/status_codes/status_codes.h>
#include <balancer/kernel/matcher/matcher.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/net/address.h>
#include <balancer/kernel/thread/threadedqueue.h>

#include <util/string/cast.h>
#include <util/system/compiler.h>

using namespace NSrvKernel;
using namespace NModPinger;

Y_TLS(pinger) {
    TSharedFileExistsChecker EnableTcpCheckChecker;
    TSharedFileReReader SwitchOffFileReader;
    THolder<TTimeLimitedAvgTracker<1000>> Stats;
    TCoroutine PingerCont;
    bool ForcedSwitchOff = false;
    bool Enabled = true;
    TSharedFileReReader::TData SwitchOffFileData;
};

MODULE_WITH_TLS(pinger) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
        , PingRequestData_("GET /robots.txt HTTP/1.1\r\nHost: yandex.ru\r\n\r\n")
        , AdminRequestUri_("/robots.txt")
        , Delay_(TDuration::Seconds(5.0))
        , HistTime_(TDuration::Seconds(10.0))
        , LowSuccRate_(0.5)
        , HiSuccRate_(0.7)
    {
        Config->ForEach(this);

        if (HistTime_ < Delay_) {
            ythrow TConfigParseError() << "histtime smaller than delay";
        }

        if (!Module_) {
            ythrow TConfigParseError() << "no module configured";
        }

        ParsePingRequest(PingRequest_);
    }

private:
    START_PARSE {
        ON_KEY ("ping_request_data", PingRequestData_) {
            return;
        }

        ON_KEY("admin_request_uri", AdminRequestUri_) {
            return;
        }

        TString adminIps;
        ON_KEY("admin_ips", adminIps) {
            AdminIpsMatcher_.Reset(new TRequestIpMatcher(adminIps));
            return;
        }

        ON_KEY("delay", Delay_) {
            return;
        }

        ON_KEY("histtime", HistTime_) {
            return;
        }

        ON_KEY("lo", LowSuccRate_) {
            if (0.0 > LowSuccRate_ || LowSuccRate_ > 1.0) {
                ythrow TConfigParseError() << "\"lo\" must be in [0.0, 1.0]";
            }
            return;
        }

        ON_KEY("hi", HiSuccRate_) {
            if (0.0 > HiSuccRate_ || HiSuccRate_ > 1.0) {
                ythrow TConfigParseError() << "\"hi\" must be in [0.0, 1.0]";
            }
            return;
        }

        ON_KEY("enable_tcp_check_file", EnableTcpCheckFile_) {
            return;
        }

        ON_KEY("switch_off_file", SwitchOffFile_) {
            return;
        }

        ON_KEY("switch_off_key", SwitchOffKey_) {
            return;
        }

        if (key == "status_codes") {
            if (!StatusCodeReactions_) {
                StatusCodeReactions_.Reset(new TStatusCodeReactions());
                StatusCodeReactions_->SetUnknownIsBad(true);
            }
            TStatusCodeReactionParser parser(*StatusCodeReactions_, TStatusCodeReaction::Good());
            value->AsSubConfig()->ForEach(&parser);
            return;
        }

        if (key == "status_codes_exceptions") {
            if (!StatusCodeReactions_) {
                StatusCodeReactions_.Reset(new TStatusCodeReactions());
                StatusCodeReactions_->SetUnknownIsBad(true);
            }
            TStatusCodeReactionParser parser(*StatusCodeReactions_, TStatusCodeReaction::Bad());
            value->AsSubConfig()->ForEach(&parser);
            return;
        }

        if (key == "admin_error_replier") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(AdminErrorReplier_);
            return;
        }

        if (key == "module" ) {
            TSubLoader(Copy(value->AsSubConfig())).Swap(Module_);
            return;
        }
    } END_PARSE

    void ParsePingRequest(TRequest& pingRequest) {
        TryRethrowError(pingRequest.Parse(PingRequestData_));
    }

    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        THolder<TTls> tls = MakeHolder<TTls>();

        if (process->WorkerType() != NProcessCore::TChildProcessType::Default) {
            return tls;
        }

        if (!!EnableTcpCheckFile_) {
            tls->EnableTcpCheckChecker =
                    process->SharedFiles()->FileChecker(EnableTcpCheckFile_, TDuration::Seconds(1));
        }

        if (!!SwitchOffFile_) {
            tls->SwitchOffFileReader =
                    process->SharedFiles()->FileReReader(SwitchOffFile_, TDuration::Seconds(1));
        }

        tls->Stats.Reset(new TTimeLimitedAvgTracker<1000>(HistTime_));

        tls->PingerCont = TCoroutine{"pinger", &process->Executor(), [&, process, tlsPtr = tls.Get()] {
            auto* const cont = process->Executor().Running();
            do {
                TRandomAddr addrOwner;
                TAddrHolder addr(&addrOwner);
                TTcpConnProps tcpConnProps(*process, addr, addr, nullptr);
                tcpConnProps.SkipKeepalive = true;
                TConnProps properties(tcpConnProps, Now(), 0, nullptr);

                TRequest request = PingRequest_;

                TNullStream nullStream;
                TAccessLogOutput nullAccessLog;
                TConnDescr descr(nullStream, nullStream, properties);
                descr.Request = &request;
                if (!StatusCodeReactions_) {
                    Y_UNUSED(RunExceptionAsError(descr, *tlsPtr));
                } else {
                    Y_UNUSED(RunStatusCodeReaction(descr, *tlsPtr));
                }
                cont->SleepT(Delay_);
            } while (!cont->Cancelled());
        }};

        return tls;
    }

    TError RunExceptionAsError(const TConnDescr& descr, TTls& tls) const {
        Y_TRY(TError, error) {
            return Module_->Run(descr);
        } Y_CATCH {
            if (error.GetAs<TForceStreamClose>()) {
                OnSuccess(tls);
            } else {
                OnFailure(tls);
            }
            return error;
        }
        OnSuccess(tls);
        return {};
    }

    TError RunStatusCodeReaction(const TConnDescr& descr, TTls& tls) const  {
        ui32 statusCode = 0;
        auto output = MakeHttpOutput([&](TResponse&& response, const bool forceClose, TInstant deadline) {
            statusCode = response.ResponseLine().StatusCode;
            return descr.Output->SendHead(std::move(response), forceClose, deadline);
        }, [&](TChunkList lst, TInstant deadline) {
            return descr.Output->Send(std::move(lst), deadline);
        }, [](THeaders&&, TInstant) {
            return TError{};
        });

        Y_TRY(TError, error) {
            return Module_->Run(descr.CopyOut(output));
        } Y_CATCH {
            if (error.GetAs<TForceStreamClose>()) {
                OnSuccess(tls, statusCode);
            } else {
                OnFailure(tls);
            }
            return error;
        }
        OnSuccess(tls, statusCode);
        return {};
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (IsAdminRequest(descr) && IsAdminIp(descr)) {
            if (!ForcedSwitchOff(tls) &&
                (tls.Enabled || tls.EnableTcpCheckChecker.Exists()))
            {
                TResponse response = BuildResponse().Version11()
                                                    .Code(HTTP_OK)
                                                    .ContentLength(0);
                Y_TRY(TError, error) {
                    Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(response), false, TInstant::Max()));
                    return descr.Output->SendEof(TInstant::Max());
                } Y_CATCH {
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "client write error");
                    return error;
                };
                descr.ExtraAccessLog << " 1";
                descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "success");
            } else {
                descr.ExtraAccessLog << " 0";
                Y_DEFER {
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "too low succ rate");
                };
                Y_REQUIRE(AdminErrorReplier_,
                          THttpError{503} << "too low succ rate " << tls.Stats->M());
                Y_PROPAGATE_ERROR(AdminErrorReplier_->Run(descr));
            }
        } else if (!StatusCodeReactions_) {
            Y_PROPAGATE_ERROR(RunExceptionAsError(descr, tls));
        } else {
            Y_PROPAGATE_ERROR(RunStatusCodeReaction(descr, tls));
        }
        return {};
    }

    bool IsAdminRequest(const TConnDescr& descr) const noexcept {
        if (Y_UNLIKELY(!descr.Request)) {
            return false;
        }

        return descr.Request->RequestLine().GetURL() == AdminRequestUri_;
    }

    bool IsAdminIp(const TConnDescr& descr) const noexcept {
        if (AdminIpsMatcher_) {
            return AdminIpsMatcher_->Match(descr);
        }
        return true;
    }

    bool ForcedSwitchOff(TTls& tls) const noexcept {
        try {
            UpdateSwitchOff(tls);
            return tls.ForcedSwitchOff;
        } catch (const yexception& e) {
        }
        return false;
    }

    void UpdateSwitchOff(TTls& tls) const {
        const auto& data = tls.SwitchOffFileReader.Data();
        if (data.Id() == tls.SwitchOffFileData.Id()) {
            return;
        }
        tls.SwitchOffFileData = data;

        class TSwitchOffReaction {
        public:
            explicit TSwitchOffReaction(const TString& switchOffKey)
                : SwitchOffKey_(switchOffKey)
            {}

            bool OnKeyValue(const TStringBuf& key, const TStringBuf& value) {
                if (key == SwitchOffKey_) {
                    HasKey_ = true;
                    KeyValue_ = FromString<double>(value) > 0;
                }
                return true;
            }

            bool HasKey() const noexcept {
                return HasKey_;
            }

            bool KeyValue() const noexcept {
                return KeyValue_;
            }

        private:
            const TString& SwitchOffKey_;
            bool HasKey_ = false;
            bool KeyValue_ = false;
        };

        NSrvKernel::TKvFileConsumer<TSwitchOffReaction> consumer(SwitchOffKey_);
        ProcessKvData(consumer, TStringBuf(tls.SwitchOffFileData.Data()));
        if (consumer.HasKey()) {
            tls.ForcedSwitchOff = consumer.KeyValue();
        } else {
            tls.ForcedSwitchOff = false;
        }
    }

    void OnSuccess(TTls& tls) const noexcept {
        tls.Stats->Add(1.0);

        if (!tls.Enabled && tls.Stats->M() > HiSuccRate_) {
            tls.Enabled = true;
        }
    }

    void OnSuccess(TTls& tls, unsigned statusCode) const noexcept {
        if (Y_UNLIKELY(!StatusCodeReactions_)) {
            OnSuccess(tls);
        }

        if (StatusCodeReactions_->IsBad(statusCode)) {
            OnFailure(tls);
        } else {
            OnSuccess(tls);
        }
    }

    void OnFailure(TTls& tls) const noexcept {
        tls.Stats->Add(0.0);

        if (tls.Enabled && tls.Stats->M() < LowSuccRate_) {
            tls.Enabled = false;
        }
    }

    bool DoExtraAccessLog() const noexcept override {
        return true;
    }

private:
    TString PingRequestData_;
    TString AdminRequestUri_;
    THolder<TRequestIpMatcher> AdminIpsMatcher_;
    TDuration Delay_;
    TDuration HistTime_;

    double LowSuccRate_ = 0.;
    double HiSuccRate_ = 0.;

    TString EnableTcpCheckFile_;

    TString SwitchOffFile_;

    TString SwitchOffKey_{ "switch_off" };

    THolder<IModule> Module_;
    THolder<IModule> AdminErrorReplier_;

    THolder<TStatusCodeReactions> StatusCodeReactions_;

    TRequest PingRequest_;
};

IModuleHandle* NModPinger::Handle() {
    return TModule::Handle();
}
