#include "module.h"

#include <balancer/modules/antirobot_common/checker_export.h>

#include <balancer/kernel/ban/banned_addresses.h>
#include <balancer/kernel/custom_io/rewind.h>
#include <balancer/kernel/fs/shared_file_exists_checker.h>
#include <balancer/kernel/http/parser/http.h>
#include <balancer/kernel/log/errorlog.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/process/main_config.h>
#include <balancer/kernel/regexp/regexp_pire.h>

#include <util/datetime/base.h>
#include <util/generic/ptr.h>
#include <util/generic/ylimits.h>

using namespace NSrvKernel;
using namespace NModAntiRobot;

namespace NModAntiRobot {

struct TAntirobotResult {
    bool IsRobot = false;
    bool SentRedirect = false;
    bool ShouldBan = false;
};

const TString XAntirobotSetCookie = "x-antirobot-set-cookie";

} // namespace NModAntiRobot

Y_TLS(antirobot) {
    TTls(bool cutRequest) : CutRequestDefault(cutRequest)
    {}

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

    bool CutRequest() const noexcept {
        if (!NoCutRequestFileChecker.Empty()) {
            return !NoCutRequestFileChecker.Exists();
        }

        return CutRequestDefault;
    }

    bool BanAddressesDisableFileExists() const noexcept {
        return BanAddressesDisableFileChecker.Exists();
    }

    bool CutRequestDefault = false;
    TSharedFileExistsChecker KillSwitchChecker;
    TSharedFileExistsChecker NoCutRequestFileChecker;
    TSharedFileExistsChecker BanAddressesDisableFileChecker;
};

MODULE_WITH_TLS_BASE(antirobot, TModuleWithSubModule) {
private:
    TAntirobotResult IsRobot(const TConnDescr& descr, TTls& tls, TVector<TStringStorage>& cookieHeader) const noexcept {
        const size_t cutRequestBytes = descr.Request && !descr.Request->Props().UpgradeRequested
            ? CutRequestBytes_ : 0;

        const bool cutRequest = descr.Request && !descr.Request->Props().UpgradeRequested
            ? tls.CutRequest() : true;

        TChecker checker(Checker_.Get(), descr, AntirobotRequest_, cutRequestBytes, false, cutRequest);

        Y_TRY(TError, error) {
            return checker.Run();
        } Y_CATCH {
            LOG_ERROR(TLOG_ERR, descr, "IsRobot check returns error: " << GetErrorMessage(error));
            return { false, false, false };
        }

        const bool isRobot = checker.IsRobot();
        const bool sentRedirect = checker.SentRedirect();
        const bool shouldBan = checker.ShouldBan();

        if (!isRobot) {
            TResponse& resp = checker.AntirobotResponse();

            for (const auto& name : ANTIROBOT_HEADERS) {
                descr.Request->Headers().Replace(name, resp.Headers().GetValuesMove(name));
            }

            cookieHeader = resp.Headers().GetValuesMove(XAntirobotSetCookie);
        }

        return TAntirobotResult{ isRobot, sentRedirect, shouldBan };
    }

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

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

        if (!Checker_) {
            ythrow TConfigParseError() << "no checker configured";
        }

        if (!CutRequest_) {
            Y_WARN_ONCE("\"cut_request\" is not set. Will fail on big inputs.");
        }

        TryRethrowError(AntirobotRequest_.Parse("POST /fullreq HTTP/1.1\r\n\r\n"));
    }

private:
    START_PARSE {
        ON_KEY("cut_request", CutRequest_) {
            return;
        }

        ON_KEY("no_cut_request_file", NoCutRequestFile_) {
            return;
        }

        ON_KEY("cut_request_bytes", CutRequestBytes_) {
            return;
        }

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

        if (key == "forward_suggest_robot") {
            return;
        }

        if (key == "checker") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(Checker_);
            return;
        } else if (key == "module") {
            TSubLoader(Copy(value->AsSubConfig())).Swap(Submodule_);
            return;
        }
    } END_PARSE

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

        if (!!NoCutRequestFile_) {
            tls->NoCutRequestFileChecker = process->SharedFiles()->FileChecker(NoCutRequestFile_, TDuration::Seconds(1));
        }

        if (Control->GetMainConfig().BanAddressesDisableFile_) {
            tls->BanAddressesDisableFileChecker = process->SharedFiles()->FileChecker(Control->GetMainConfig().BanAddressesDisableFile_, TDuration::Seconds(1));
        }
        return tls;
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        if (!tls.KillSwitchFileExists()) {
            // TODO(velavokr): BALANCER-2182 - unlimited input!!!
            TLimitedRewindableInput savedInput{ *descr.Input, 1024 * 1024 * 30 };

            TVector<TStringStorage> cookieHeader;
            auto output = MakeHttpOutput([&](TResponse&& response, const bool forceClose, TInstant deadline) {
                response.Headers().Add("Set-Cookie", std::move(cookieHeader));
                return descr.Output->SendHead(std::move(response), forceClose, deadline);
            }, [&](TChunkList lst, TInstant deadline) {
                return descr.Output->Send(std::move(lst), deadline);
            }, [&](THeaders&& trailers, TInstant deadline) {
                return descr.Output->SendTrailers(std::move(trailers), deadline);
            });

            LOG_ERROR(TLOG_INFO, descr, "antirobot started");

            TConnDescr newDescr = descr.Copy(savedInput, output);

            const auto result = IsRobot(newDescr, tls, cookieHeader);
            savedInput.Rewind(); // rewind after check
            if (!descr.Request || descr.Request->Props().UpgradeRequested) {
                savedInput.ResetRewind();
            }

            const bool shouldBan = result.ShouldBan && !tls.BanAddressesDisableFileExists();
            if (!shouldBan && result.ShouldBan) {
                LOG_ERROR(TLOG_ERR, descr, "did not ban ip " << descr.RemoteAddrStr() << ": disabled by ITS");
            }

            if (!result.IsRobot && !shouldBan) {
                LOG_ERROR(TLOG_INFO, descr, "not robot request");

                const TExtraAccessLogEntry logEntry(descr, "sub_search");
                Y_PROPAGATE_ERROR(Submodule_->Run(newDescr));
            } else {
                if (result.IsRobot) {
                    LOG_ERROR(TLOG_ERR, descr, "robot request");
                }
                if (shouldBan) {
                    bool banned = Control->BannedAddresses().Add(
                        descr.RemoteAddr(),
                        Now() + TDuration::Minutes(5) + TDuration::Seconds(RandomNumber<ui64>(300)),
                        Control->GetMainConfig().BanAddressesMaxCount_
                    );
                    if (banned) {
                        LOG_ERROR(TLOG_ERR, descr, "ban ip " << descr.RemoteAddrStr());
                    } else {
                        LOG_ERROR(TLOG_ERR, descr, "could not ban ip " << descr.RemoteAddrStr() << ": limit exceeded");
                    }
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "ban ip");
                    // drop client connection
                    return Y_MAKE_ERROR(TBanAddressError{} << "ban ip by antirobot request");
                }
                if (result.SentRedirect) {
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "robot request");
                    return {};
                } else {
                    descr.ExtraAccessLog << " robot and no response";
                    descr.ExtraAccessLog.SetSummary(GetHandle()->Name(), "robot and no response");
                    return Y_MAKE_ERROR(yexception{} << "is robot and no redirect response for it");
                }
            }
        } else {
            for (const auto& name : ANTIROBOT_HEADERS) {
                descr.Request->Headers().Delete(name);
            }

            Y_PROPAGATE_ERROR(Submodule_->Run(descr));
        }
        return {};
    }

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

    void DoCheckConstraints() const override {
        if (!CheckParents([&](TStringBuf name) { return "h100" == name; })) {
            Y_WARN_ONCE("must have h100 as a parent");
        }
    }

private:
    THolder<IModule> Checker_;

    TString KillSwitchFile_;
    TString NoCutRequestFile_;

    size_t CutRequestBytes_ = Max<size_t>();
    bool CutRequest_ = false;

    TRequest AntirobotRequest_;
};

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