#include "module.h"

#include <balancer/kernel/cookie/cookie.h>
#include <balancer/kernel/fs/shared_file_exists_checker.h>
#include <balancer/kernel/fs/shared_files.h>
#include <balancer/kernel/http/parser/common_headers.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/string.h>
#include <util/random/random.h>


using namespace NConfig;
using namespace NSrvKernel;

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

    TSharedFileExistsChecker KillSwitchChecker;
};

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

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

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

    START_PARSE {
        ON_KEY("cookie", CookieName_) {
            return;
        }

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

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

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

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

    TMaybe<TRequestHash> CalcHash(const TConnDescr& descr) const noexcept {
        TMurmurHash2A<TRequestHash> hash;
        TMaybe<TRequestHash> res;

        if (auto it = descr.Request->Headers().FindValues("cookie"); it != descr.Request->Headers().end()) {
            for (auto& headerValue: it->second) {
                EveryCookieKV([&](TNameValue nv) noexcept {
                    if (nv.Name == CookieName_) {
                        hash.Update(nv.Value.data(), nv.Value.size());
                        res = 0;
                    }
                }, headerValue.AsStringBuf());
            }
        }

        if (res) {
            res = hash.Value();
        } else if (RandomizeEmptyMatch_) {
            res = RandomNumber<TRequestHash>();
        }

        return res;
    }

    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 (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 CookieName_;
    TString KillSwitchFile_;
    bool RandomizeEmptyMatch_ = true;
    bool CombineHashes_ = false;
};

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