#pragma once

#include "template.h"
#include "signature.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/iface.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/net/address.h>

namespace NModAntiDDOS {
    // This class contains info necessary to do redirection:
    // connection description and possibly customized target.
    // If redirection method finds that target finishes with
    // antiddos signature it can cut it off so that it is not
    // repeated several times.

    using namespace NSrvKernel;

    extern const TString DEFAULT_REDIRECTION_SALT;

    class TRedirectionInfo {
    public:
        TRedirectionInfo(const TConnDescr& descr) noexcept
            : Descr_(descr)
        {}

        const TConnDescr& ConnDescr() const noexcept {
            return Descr_;
        }

        TString& Target() noexcept {
            if (!Target_) {
                auto tmp = ToString(Descr_.Request->RequestLine().GetURL());
                auto it = std::find(tmp.cbegin(), tmp.cend(), ' ');
                if (it != tmp.end()) {
                    ++it;
                }
                auto itNext = std::find(it, tmp.cend(), ' ');
                tmp.erase(std::move(it, itNext, Target_->begin()), Target_->cend());
                tmp.resize(itNext - it);
                Target_ = std::move(tmp);
            }

            return *Target_;
        }

    private:
        TMaybe<TString> Target_;
        const TConnDescr& Descr_;
    };

    class TRedirectionMethod {
    public:
        TRedirectionMethod(const TString& salt) noexcept
            : Salt_(salt)
        {
        }

        virtual ~TRedirectionMethod() = default;

        virtual TError Redirect(TRedirectionInfo& info) const noexcept = 0;

        virtual TErrorOr<EProtectionLevel> GetRequestProtectionLevel(
            TRedirectionInfo& info) const noexcept
        {
            TSignatureInfo signature;
            bool loaded;
            Y_PROPAGATE_ERROR(signature.Load(info.Target()).AssignTo(loaded));
            if (!loaded) {
                return PROTECTION_BYPASS;
            }
            TSignatureInfo::CutOffSignature(info.Target());
            TString remoteHost = info.ConnDescr().RemoteAddrStr();
            bool isCorrect = signature.Verify(info.Target(), remoteHost, info.ConnDescr().Properties->Start, Salt_);
            return isCorrect ? signature.ProtectionLevel() : PROTECTION_BYPASS;
        }

        virtual double GetRateThreshold() const noexcept = 0;
        virtual EProtectionLevel GetMethodProtectionLevel() const noexcept = 0;

    protected:
        const TString& Salt_;
    };

    class TBypassRedirection: public TRedirectionMethod {
    public:
        TBypassRedirection() noexcept
            : TRedirectionMethod(DEFAULT_REDIRECTION_SALT)
        {
        }

        TError Redirect(TRedirectionInfo&) const noexcept override {
            return Y_MAKE_ERROR(yexception{} << "bypass method won't redirect");
        }

        double GetRateThreshold() const noexcept override {
            return 0;
        }

        EProtectionLevel GetMethodProtectionLevel() const noexcept override {
            return PROTECTION_BYPASS;
        }

        TErrorOr<EProtectionLevel> GetRequestProtectionLevel(
            TRedirectionInfo&) const noexcept override
        {
            return PROTECTION_DROP;
        }
    };

    class TDropRedirection: public TRedirectionMethod, public TModuleParams, public NConfig::IConfig::IFunc {
    public:
        TDropRedirection(const TModuleParams& mp)
            : TRedirectionMethod(DEFAULT_REDIRECTION_SALT)
            , TModuleParams(mp)
        {
            Config->ForEach(this);
        }

        START_PARSE {
            ON_KEY("threshold", Threshold_) {
                return;
            }
        } END_PARSE

        TError Redirect(TRedirectionInfo& redirectionInfo) const noexcept override {
            TString NAME = "antiddos"; // FIXME
            LOG_ERROR(TLOG_ERR, redirectionInfo.ConnDescr(), "applying drop redirection");
            return Y_MAKE_ERROR(THttpError{509});
        }

        double GetRateThreshold() const noexcept override {
            return Threshold_;
        }

        EProtectionLevel GetMethodProtectionLevel() const noexcept override {
            return PROTECTION_DROP;
        }

        TErrorOr<EProtectionLevel> GetRequestProtectionLevel(
            TRedirectionInfo&) const noexcept override
        {
            return PROTECTION_BYPASS;
        }

    protected:
        double Threshold_;
    };

    class TTemplateRedirection: public TRedirectionMethod, public TModuleParams, public NConfig::IConfig::IFunc {
    public:
        TTemplateRedirection(const TModuleParams& mp, const TString& salt, EProtectionLevel protLevel, TDuration validTime)
            : TRedirectionMethod(salt)
            , TModuleParams(mp)
            , ProtectionLevel_(protLevel)
            , ValidityTime_(validTime)
        {
            Config->ForEach(this);
        }

        START_PARSE {
            ON_KEY("threshold", Threshold_) {
                return;
            }

            if (key == "template") {
                Template_.Load(value->AsString());
                return;
            }
        } END_PARSE

        TError Redirect(TRedirectionInfo& info) const noexcept override {
            TString NAME = "antiddos"; // FIXME

            TString stringOut;
            TStringOutput out{ stringOut };

            TString remoteHost = info.ConnDescr().RemoteAddrStr();
            TSignedTarget target{ info.Target(), remoteHost, GetMethodProtectionLevel(),
                                  info.ConnDescr().Properties->Start + ValidityTime_, Salt_ };

            if (GetMethodProtectionLevel() == PROTECTION_HTTP_REDIRECT) {
                LOG_ERROR(TLOG_ERR, info.ConnDescr(), "applying http redirection");
                out << "HTTP/1.1 302 Moved\r\n"
                    << "Location: " << target << "\r\n";
            } else {
                LOG_ERROR(TLOG_ERR, info.ConnDescr(), "applying html redirection");
                out << "HTTP/1.1 200 OK\r\n";
            }

            out << "Content-Length: " << Template_.Length(target)
                << "\r\n"
                << "Content-Type: text/html; charset=UTF-8\r\n"
                << "\r\n";

            Template_.Write(target, out);
            return info.ConnDescr().Output->Send(TChunkList(std::move(stringOut)), TInstant::Max());
        }

        double GetRateThreshold() const noexcept override {
            return Threshold_;
        }

        EProtectionLevel GetMethodProtectionLevel() const noexcept override {
            return ProtectionLevel_;
        }

    protected:
        TRedirectionTemplate Template_;
        double Threshold_;
        EProtectionLevel ProtectionLevel_;
        TDuration ValidityTime_;
    };
}
