#include "module.h"

#include <balancer/kernel/fs/watched_state.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/http/parser/response_builder.h>
#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/memory/chunks.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/log/errorlog.h>

#include <library/cpp/http/misc/httpcodes.h>

using namespace NSrvKernel;
using namespace NModActiveCheckReply;

namespace NModActiveCheckReply {
    const size_t MIN_WEIGHT = 1;
    const size_t MAX_WEIGHT = 1000;

    struct TWeightData {
        size_t IntegerWeight;
        TString Weight;
        TString ContentString;
    };

    TWeightData MakeWeightData(size_t weight) noexcept {
        return {weight, ToString(weight), TString("rs_weight=") + ToString(weight)};
    }

    struct TOptions {
        size_t DefaultWeight = 1;
        TString WeightFile;

        TMaybe<bool> UseBody;
        TMaybe<bool> UseHeader;

        bool ZeroWeightAtShutdown = false;
    };

    struct TWeightStateWithOptions {
        TWeightStateWithOptions() = default;

        TWeightStateWithOptions(TOptions options, TSharedFiles& sharedFiles)
            : Options(std::move(options))
            , State(MakeWeightData(Options.DefaultWeight), Options.WeightFile, sharedFiles)
        {}

        TOptions Options;
        TWatchedState<TWeightData> State = MakeWeightData(Options.DefaultWeight);
    };
}

Y_TLS(active_check_reply) {
    TWeightStateWithOptions WeightStateWithOptions;

    TResponse Response;

    TMaybe<TSharedCounter> StatsWeightCounter;
};

MODULE_WITH_TLS(active_check_reply) {
public:
    explicit TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

        if (!Options_.UseHeader.Defined() && !Options_.UseBody.Defined()) {
            Options_.UseHeader = true;
            Options_.UseBody = false;
        }

        if (!Options_.UseHeader.Defined()) {
            Options_.UseHeader = false;
        }

        if (!Options_.UseBody.Defined()) {
            Options_.UseBody = false;
        }

        if (!*Options_.UseBody && !*Options_.UseHeader) {
            ythrow TConfigParseError() << "active_check_reply use_header/use_body missconfiguration";
        }

        TStringStream statName("active_check_reply-");
        if (Uuid_) {
            statName << Uuid_ << '-';
        }
        statName << "weight";
        SharedCounter_ = Control->SharedStatsManager().MakeGauge(statName.Str()).AllowDuplicate().Aggregation(TWorkerAggregation::Avg).Build();

    }

private:
    THolder<TTls> DoInitTls(IWorkerCtl* worker) override {
        auto tls = MakeHolder<TTls>();

        tls->Response = (ForceConnClose_ ? BuildResponse().Version10() : BuildResponse().Version11()).Code(HTTP_OK).ContentLength(0);
        tls->StatsWeightCounter = TSharedCounter(*SharedCounter_, worker->WorkerId());

        tls->WeightStateWithOptions = TWeightStateWithOptions(Options_, *worker->SharedFiles());

        return tls;
    }

    START_PARSE {
        ON_KEY("uuid", Uuid_) {
            return;
        }

        ON_KEY("default_weight", Options_.DefaultWeight) {
            if (Options_.DefaultWeight < MIN_WEIGHT || Options_.DefaultWeight > MAX_WEIGHT) {
                ythrow TConfigParseError() << "active_check_reply default_weight violates constraints";
            }
            return;
        }

        ON_KEY("weight_file",  Options_.WeightFile) {
            return;
        }

        ON_KEY("use_header", Options_.UseHeader) {
            return;
        }

        ON_KEY("use_body", Options_.UseBody) {
            return;
        }

        ON_KEY("zero_weight_at_shutdown", Options_.ZeroWeightAtShutdown) {
            return;
        }

        ON_KEY("push_checker_address", PushCheckerAddress_) {
            return;
        }

        ON_KEY("force_conn_close", ForceConnClose_) {
            return;
        }

        if (key == "subnet_filter") {
            return;
        }

        if (key == "subnet_filter_disable_file") {
            return;
        }

        if (key == "use_dynamic_weight") {
            return;
        }

        if (key == "dynamic_weight_disable_file") {
            return;
        }
    } END_PARSE

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (PushCheckerAddress_ && descr.CpuLimiter() != nullptr) {
            descr.CpuLimiter()->PushCheckerAddress(descr.RemoteAddr());
        }

        const bool zeroWeight = Options_.ZeroWeightAtShutdown && descr.Process().IsShuttingDown() && descr.Process().BlockNewRequests();
        const TWeightData& weightData = zeroWeight ? MakeWeightData(0) : GetWeightData(descr, tls);
        tls.StatsWeightCounter->Set(weightData.IntegerWeight);

        TResponse response = tls.Response;

        if (*Options_.UseHeader) {
            response.Headers().Add("RS-Weight", weightData.Weight);
        }

        const bool haveBody = *Options_.UseBody && !weightData.ContentString.empty();

        if (haveBody) {
            response.Props().ContentLength = weightData.ContentString.size();
        }

        Y_TRY(TError, error) {
            Y_PROPAGATE_ERROR(descr.Output->SendHead(std::move(response), ForceConnClose_, TInstant::Max()));
            if (haveBody) {
                Y_PROPAGATE_ERROR(descr.Output->Send(TChunkList{weightData.ContentString}, TInstant::Max()));
            }
            Y_PROPAGATE_ERROR(descr.Output->SendEof(TInstant::Max()));

            descr.ExtraAccessLog << ' ' << weightData.ContentString;
            return SkipAll(descr.Input, TInstant::Max());
        } Y_CATCH {
            descr.ExtraAccessLog << " error";
            descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "client write error");
            return error;
        }
        descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "success");
        return {};
    }

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

    const TWeightData& GetWeightData(const TConnDescr& descr, TTls& tls) const noexcept {
        const auto& options = tls.WeightStateWithOptions.Options;

        return tls.WeightStateWithOptions.State.Get(
            [&options, &descr](const TString& data, TWeightData& currentV, const TWeightData& defaultV) {
                TString firstLine;
                if (TStringInput(data).ReadLine(firstLine)) {
                    size_t weight = 0;
                    if (TryFromString(firstLine, weight)) {
                        weight = ClampVal(weight, 0ul, MAX_WEIGHT);
                        currentV = MakeWeightData(weight);
                        return;
                    } else {
                        LOG_ERROR(TLOG_ERR, descr, "can't parse weight from "
                            << options.WeightFile << ": " << firstLine);
                    }
                }

                currentV = defaultV;
            }
        );
    }

private:
    TString Uuid_;
    bool PushCheckerAddress_ = false;
    bool ForceConnClose_ = true;

    TOptions Options_;

    TMaybe<TSharedCounter> SharedCounter_;
};

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