#include "module.h"
#include "redirection.h"
#include "ratecounter.h"

#include <balancer/kernel/custom_io/stream.h>
#include <balancer/kernel/module/module.h>

#include <util/generic/map.h>
#include <util/generic/ptr.h>
#include <util/stream/str.h>

using namespace NModAntiDDOS;
using namespace NConfig;
using namespace NSrvKernel;

namespace {
    static const char* LEVEL_NAMES[] = {
        "bypass",
        "http_redirect",
        "html_redirect",
        "js_redirect",
        //"show_captcha",
        "drop"
    };

    TString Verdict(EProtectionLevel level) {
        return LEVEL_NAMES[level];
    }
}

MODULE_BASE(antiddos, TModuleWithSubModule) {
private:
    using TRateToProtLevelMap = TMap<double, EProtectionLevel>;
    using TProtLevelToRedirMap = TMap<EProtectionLevel, THolder<TRedirectionMethod>>;

public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
    {
        Zero(RedirectionCounter_);

        Config->ForEach(this);

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

        RateCounter_.Init(RateToProtLevel_.size() > 0 ? (--RateToProtLevel_.end())->first : 0);
    }

private:
    START_PARSE {
        ON_KEY("salt", Salt_) {
            return;
        }

        THolder<TRedirectionMethod> meth;

        if (key == LEVEL_NAMES[PROTECTION_HTTP_REDIRECT]) {
            meth.Reset(new TTemplateRedirection(Copy(value->AsSubConfig()), Salt_, PROTECTION_HTTP_REDIRECT, SignValidTime));
        } else if (key == LEVEL_NAMES[PROTECTION_HTML_REDIRECT]) {
            meth.Reset(new TTemplateRedirection(Copy(value->AsSubConfig()), Salt_, PROTECTION_HTML_REDIRECT, SignValidTime));
        } else if (key == LEVEL_NAMES[PROTECTION_JAVASCRIPT_REDIRECT]) {
            meth.Reset(new TTemplateRedirection(Copy(value->AsSubConfig()), Salt_, PROTECTION_JAVASCRIPT_REDIRECT, SignValidTime));
        } else if (key == LEVEL_NAMES[PROTECTION_DROP]) {
            meth.Reset(new TDropRedirection(Copy(value->AsSubConfig())));
        } else {
            Submodule_.Reset(Loader->MustLoad(key, Copy(value->AsSubConfig())).Release());
            meth.Reset(new TBypassRedirection());
        }

        if (!!meth) {
            RateToProtLevel_[meth->GetRateThreshold()] = meth->GetMethodProtectionLevel();
            ProtLevelToRedirMethod_[meth->GetMethodProtectionLevel()] = std::move(meth);

            return;
        }
    } END_PARSE

    EProtectionLevel GetProtectionLevel(double rate) const noexcept {
        return (--(RateToProtLevel_.upper_bound(rate)))->second;
    }

    TError DoRun(const TConnDescr& descr) const noexcept override {
        EProtectionLevel currentLevel = GetProtectionLevel(RateCounter_.CalculateRate(descr.Properties->Start));
        const TRedirectionMethod* redirMethod = ProtLevelToRedirMethod_[currentLevel].Get();
        TRedirectionInfo redirInfo{ descr };
        EProtectionLevel requestLevel;
        if (TError error = redirMethod->GetRequestProtectionLevel(redirInfo).AssignTo(requestLevel)) {
            descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "signature error");
            return error;
        }

        if (requestLevel >= currentLevel) {
            ++RedirectionCounter_[PROTECTION_BYPASS];
            RateCounter_.AddRequest(descr.Properties->Start);
            Y_PROPAGATE_ERROR(Submodule_->Run(descr));
        } else {
            ++RedirectionCounter_[currentLevel];
            TError error = redirMethod->Redirect(redirInfo);
            descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), Verdict(currentLevel));
            return error;
        }
        return {};
    }

private:
    mutable TRateCounter RateCounter_;
    mutable TRateToProtLevelMap RateToProtLevel_;
    mutable TProtLevelToRedirMap ProtLevelToRedirMethod_;
    TString Salt_;

    mutable ui64 RedirectionCounter_[PROTECTION_LEVELS_NUMBER];

    static const TDuration SignValidTime;
};

const TDuration TModule::SignValidTime = TDuration::Minutes(10);

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

