#include "module.h"

#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/fs/shared_file_exists_checker.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/memory/chunks.h>
#include <balancer/kernel/module/iface.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/regexp/regexp_pire.h>

#include <util/generic/ptr.h>
#include <util/generic/strbuf.h>
#include <util/generic/string.h>
#include <util/random/random.h>


using namespace NConfig;
using namespace NSrvKernel;

Y_TLS(headers_hasher) {
    bool KillSwitchFileExists() const noexcept {
        return KillSwitchChecker.Exists();
    }

    TSharedFileExistsChecker KillSwitchChecker;
};


MODULE_WITH_TLS_BASE(headers_hasher, TModuleWithSubModule) {
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

        if (!Submodule_) {
            ythrow TConfigParseError() << "no submodule configured";
        }

        if (!HeaderName_) {
            ythrow TConfigParseError() << "\"header_name\" must not be empty";
        }

        HeaderNameFsm_.Reset(new TFsm(HeaderName_, TFsm::TOptions().SetCaseInsensitive(true).SetSurround(Surround_)));
    }

    START_PARSE {
        ON_KEY("header_name", HeaderName_) {
            return;
        }

        ON_KEY("surround", Surround_) {
            return;
        }

        ON_KEY("randomize_empty_match", RandomizeEmptyMatch_) {
            return;
        }

        ON_KEY("combine_hashes", CombineHashes_) {
            return;
        }

        ON_KEY("file_switch", KillSwitchFile_) {
            return;
        }

        {
            Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
            return;
        }
    } END_PARSE


    TMaybe<ui64> CalcHash(const TConnDescr& descr) const noexcept {
        const TRequest* request = descr.Request;
        TMurmurHash2A<TRequestHash> hash;
        bool found = false;

        for (const auto& header: request->Headers()) {
            if (Match(*HeaderNameFsm_, header.first.AsStringBuf())) {
                for (const auto& headerValue : header.second) {
                    found = found || !headerValue.Empty();
                    auto headerValueBuffer = headerValue.AsStringBuf();
                    hash.Update(headerValueBuffer.Data(), headerValueBuffer.Size());
                }
            }
        }

        if (found) {
            return hash.Value();
        } else {
            return Nothing();
        }
    }

    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        auto tls = MakeHolder<TTls>();
        if (!!KillSwitchFile_) {
            tls->KillSwitchChecker = process->SharedFiles()->FileChecker(KillSwitchFile_, TDuration::Seconds(1));
        }
        return tls;
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (!tls.KillSwitchFileExists()) {
            auto result = CalcHash(descr);
            if (!result && RandomizeEmptyMatch_) {
                result = RandomNumber<NSrvKernel::TRequestHash>();
            }
            if (CombineHashes_ && descr.Hash) {
                result = result ? CombineHashes(*result, descr.Hash) : descr.Hash;
            }
            Y_PROPAGATE_ERROR(Submodule_->Run(descr.Copy(result.GetOrElse(0))));
        } else {
            Y_PROPAGATE_ERROR(Submodule_->Run(descr));
        }
        return {};
    }

private:
    TString HeaderName_;
    TString KillSwitchFile_;
    THolder<TFsm> HeaderNameFsm_;
    bool Surround_ = false;
    bool RandomizeEmptyMatch_ = true;
    bool CombineHashes_ = false;
};

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