#include "module.h"

#include <extsearch/images/robot/library/thumbid/thumbid.h>

#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/log/errorlog.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/regexp/regexp_re2.h>

#include <library/cpp/config/sax.h>

#include <util/generic/strbuf.h>
#include <util/generic/vector.h>
#include <util/string/hex.h>

using namespace NConfig;
using namespace NSrvKernel;

using TModulePtr = TSimpleSharedPtr<IModule>;
using TModuleData = std::pair<size_t, TModulePtr>;

namespace {
    bool ModulesCmp(const TModuleData &x, const TModuleData &y) {
        return x.first < y.first;
    }
}

MODULE(thumb_consistent_hash) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Config->ForEach(this);

        if (!Regexp) {
            ythrow TConfigParseError() << "no regexp configured";
        }

        if (Modules.empty()) {
            ythrow TConfigParseError() << "no modules configured";
        }

        Sort(Modules.begin(), Modules.end(), ModulesCmp);
        // Verify presence of all shards
        for (size_t i = 0; i < Modules.size(); ++i) {
            if (Modules[i].first != i) {
                ythrow yexception() << "Module number " << Modules[i].first << " at position " << i;
            }
            if (!Modules[i].second) {
                ythrow yexception() << "Module number " << i << " not initialized";
            }
        }
    }

private:
    START_PARSE {
        TString idRegExp;
        ON_KEY ("id_regexp", idRegExp) {
            Regexp.Reset(new TRegexp(idRegExp));
            return;
        }

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

        Modules.push_back(std::make_pair(FromString<size_t>(key), TSubLoader(Copy(value->AsSubConfig())).Release()));

        return;

    } END_PARSE

    TError DoRun(const TConnDescr& descr) const noexcept override {
        LOG_ERROR(TLOG_INFO, descr, "thumb_consistent_hash started");

        TVector<TStringBuf> capturedGroups;

        if (Regexp->Extract(descr.Request->RequestLine().CGI.AsStringBuf(), &capturedGroups, true) &&
            capturedGroups.size() > 1) {
            NThDb::TThumbId id;
            TStringBuf idStr = capturedGroups[1];
            Y_TRY(TError, error) {
                try {
                    id = FromString<NThDb::TThumbId>(idStr);
                } Y_TRY_STORE(TFromStringException, yexception);
                return {};
            } Y_CATCH {
                LOG_ERROR(TLOG_ERR, descr, GetErrorMessage(error));
                if (!Default) {
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "no default");
                    return error;
                } else {
                    return Default->Run(descr);
                }
            }

            size_t shard = ConsistentHashing(id, Modules.size()); // will not throw
            descr.ExtraAccessLog << " shard " << shard;

            IModule* module = Modules[shard].second.Get();
            Y_PROPAGATE_ERROR(module->Run(descr));
        } else {
            Y_PROPAGATE_ERROR(InvokeDefault(descr));
        }
        return {};
    }

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

    TError InvokeDefault(const TConnDescr &descr) const {
        if (!Default) {
            descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "no default");
            return {};
        }
        return Default->Run(descr);
    }

private:
    TVector<TModuleData> Modules;
    THolder<TRegexp> Regexp;
    THolder<IModule> Default;
};

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