#include "module.h"

#include "balancer/kernel/helpers/default_instance.h"
#include <balancer/kernel/helpers/misc.h>
#include <balancer/kernel/http/parser/common_headers.h>
#include <balancer/kernel/http/parser/header_validation.h>
#include <balancer/kernel/log/errorlog.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/process/thread_info.h>
#include <balancer/kernel/requester/requester.h>

#include <laas/lib/ip_properties/reader/reader.h>

#include <util/thread/singleton.h>


using namespace NConfig;
using namespace NSrvKernel;
using namespace NRegExp;


struct TRestrictedHeadersFsm : public TFsm , public TWithDefaultInstance<TRestrictedHeadersFsm> {
    TRestrictedHeadersFsm() noexcept
        : TFsm("content-length|transfer-encoding|connection", TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

struct TRegionHeaderFsm : public TFsm , public TWithDefaultInstance<TRegionHeaderFsm> {
    TRegionHeaderFsm() noexcept
        : TFsm("X-Region-.*", TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

struct TIpPropertiesHeaderFsm : public TFsm , public TWithDefaultInstance<TIpPropertiesHeaderFsm> {
    TIpPropertiesHeaderFsm() noexcept
        : TFsm("X-IP-Properties", TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

constexpr TStringBuf XIpPropertiesCountry = "X-IP-Properties-Country";

struct TIpPropertiesCountryHeaderFsm : public TFsm , public TWithDefaultInstance<TIpPropertiesHeaderFsm> {
    TIpPropertiesCountryHeaderFsm() noexcept
        : TFsm(XIpPropertiesCountry, TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

struct TXIsTouristHeaderFsm: public TFsm, public TWithDefaultInstance<TXIsTouristHeaderFsm> {
    TXIsTouristHeaderFsm() noexcept
        : TFsm("X-Is-Tourist", TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

struct TLaasHeaderFsm : public TFsm , public TWithDefaultInstance<TLaasHeaderFsm> {
    TLaasHeaderFsm() noexcept
        : TFsm(
            TRegionHeaderFsm::Instance()
            | TIpPropertiesHeaderFsm::Instance()
            | TIpPropertiesCountryHeaderFsm::Instance()
            | TXIsTouristHeaderFsm::Instance()
        )
    {}
};

struct TXForwardedForYFsm : public TFsm , public TWithDefaultInstance<TXForwardedForYFsm> {
    TXForwardedForYFsm() noexcept
        : TFsm("X-Forwarded-For-Y", TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

struct TXUrlPrefixFsm : public TFsm , public TWithDefaultInstance<TXUrlPrefixFsm> {
    TXUrlPrefixFsm() noexcept
        : TFsm("X-Url-Prefix", TFsm::TOptions().SetCaseInsensitive(true))
    {}
};

class TParamsCollector {
public:
    TParamsCollector(const TConnDescr& descr, const TRequestLine& geoRequest, const TString& geoHost, const TFsm& fsm) noexcept
        : Descr_(descr)
        , RequestLine_(geoRequest)
        , GeoHost_(geoHost)
        , Fsm_(fsm)
    {}

    TError Process(bool& processedOriginalRequest) noexcept {
        if (GeoRequest_.Get()) { // already processed it succesfully
            processedOriginalRequest = true;
            return {};
        }

        auto geoRequest = MakeHolder<TRequest>(RequestLine_);
        Y_PROPAGATE_ERROR(ProcessRequests(*geoRequest, processedOriginalRequest));
        GeoRequest_.Swap(geoRequest);
        return {};
    }

    bool Ok() const noexcept {
        return !!GeoRequest_;
    }

    bool IpHeader() const noexcept {
        return IpHeader_;
    }

    TRequest&& GeoRequest() noexcept {
        return std::move(*GeoRequest_);
    }

private:
    TError ProcessRequests(TRequest& geoRequest, bool& processedOriginalRequest) noexcept {
        Y_REQUIRE(Descr_.Request, // TODO: maybe logging
                  yexception{} << "no request in descr getting geobase response");

        return DoProcessRequests(geoRequest, *Descr_.Request, processedOriginalRequest);
    }

    TError DoProcessRequests(TRequest& geoRequest, TRequest& originalRequest,
                             bool& processedOriginalRequest) noexcept
    {
        TString urlPrefix = Descr_.Properties->UserConnIsSsl ? "https://" : "http://";
        TStringStorage* urlPrefixHeaderValue = nullptr;
        TStringBuf ipHeaderKey;
        TStringStorage ipHeaderValue;
        bool hadHost = false;

        for (auto it = originalRequest.Headers().begin(); it != originalRequest.Headers().end();) {
            auto& header = *it;
            TMatcher matcher{Fsm_};

            if (Match(matcher, header.first.AsStringBuf()).Final()) {
                switch (*matcher.MatchedRegexps().first) {
                case 0: // ipFsm
                    IpHeader_ = true;
                    if (!ipHeaderKey) {
                        ipHeaderKey = header.first.AsStringBuf();
                        ipHeaderValue = header.second[0];
                    }
                    geoRequest.Headers().Add(header.first.AsString(), THeaders::MakeOwned(header.second));
                    break;
                case 1: // Host
                    if (!hadHost) {
                        hadHost = true;
                        urlPrefix += header.second[0].AsStringBuf();
                        geoRequest.Headers().Add("Host", GeoHost_);
                    }
                    break;
                case 2: // urlPrefix
                    if (!urlPrefixHeaderValue) {
                        geoRequest.Headers().Add(header.first.AsString(), header.second[0].AsString());
                        urlPrefixHeaderValue = &geoRequest.Headers().GetValuesRef(header.first.AsStringBuf())[0];
                    }
                    break;
                case 3: // restricted
                    break;
                case 4: // LaaS answer header
                case 5: // X-Region headers
                case 6: // X-IP-Properties header
                case 7: // X-Is-Tourist header
                     originalRequest.Headers().erase(it++);

                    continue;
                default:
                    geoRequest.Headers().Add(header.first.AsString(), THeaders::MakeOwned(header.second));
                    break;
                }
            } else {
                geoRequest.Headers().Add(header.first.AsString(), THeaders::MakeOwned(header.second));
            }
            ++it;
        }
        processedOriginalRequest = true;

        Y_REQUIRE(ipHeaderValue, yexception{} << "no ip header found");

        if (!hadHost) {
            urlPrefix += GeoHost_;
            geoRequest.Headers().Add("Host", GeoHost_);
        }

        if (!TCIOps()(ipHeaderKey, "x-forwarded-for-y")) {
            geoRequest.Headers().Replace("X-Forwarded-For-Y", ipHeaderValue.AsString());
        }

        urlPrefix += originalRequest.RequestLine().Path.AsStringBuf();

        if (urlPrefixHeaderValue) {
            *urlPrefixHeaderValue = TStringStorage(std::move(urlPrefix));
        } else {
            geoRequest.Headers().Add("X-Url-Prefix", std::move(urlPrefix));
        }

        return {};
    }

private: // input
    const TConnDescr& Descr_;
    const TRequestLine& RequestLine_;
    const TString& GeoHost_;
    const TFsm& Fsm_;

private: // output
    THolder<TRequest> GeoRequest_;
    bool IpHeader_ = false;
};


bool CopyRegionHeaders(THeaders& from, THeaders& to) noexcept {
    bool found = false;
    for (auto& hdr : from) {
        if (Match(TLaasHeaderFsm::Instance(), hdr.first.AsStringBuf())) {
            if (Match(TIpPropertiesHeaderFsm::Instance(), hdr.first.AsStringBuf())) {
                for (const auto& hdrVal : hdr.second) {
                    NLaas::TIpProperties props;
                    if (NLaas::TryParseFromBase64(hdrVal.AsStringBuf(), props)) {
                        to.Add(XIpPropertiesCountry, ToString(props.countryidbyip()));
                    }
                }
            }
            to.Add(hdr.first.AsString(), THeaders::MakeOwned(std::move(hdr.second)));
            found = true;
        }
    }

    return found;
}

namespace {
struct TSharedCounters {
    explicit TSharedCounters(TSharedStatsManager& statsManager)
        : LaasNoAnswer(statsManager.MakeCounter("geobase-laas_no_answer").AllowDuplicate().Build())
        , LaasNoHeader(statsManager.MakeCounter("geobase-laas_no_header").AllowDuplicate().Build())
        , LocalFails(statsManager.MakeCounter("geobase-geobase_fails").AllowDuplicate().Build())
        , NoIpInRequest(statsManager.MakeCounter("geobase-no_ip").AllowDuplicate().Build())
        , TrustedRequests(statsManager.MakeCounter("geobase-trusted_reqs").AllowDuplicate().Build())
    {}

    TSharedCounter LaasNoAnswer;
    TSharedCounter LaasNoHeader;
    TSharedCounter LocalFails;
    TSharedCounter NoIpInRequest;
    TSharedCounter TrustedRequests;
};
}

Y_TLS(geobase) {
    TTls(TSharedCounters& holders, size_t workerId)
        : LaasNoAnswer(holders.LaasNoAnswer, workerId)
        , LaasNoHeader(holders.LaasNoHeader, workerId)
        , LocalFails(holders.LocalFails, workerId)
        , NoIpInRequest(holders.NoIpInRequest, workerId)
        , TrustedRequests(holders.TrustedRequests, workerId)
    {}

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

    TSharedFileExistsChecker KillSwitchChecker;

    TSharedCounter LaasNoAnswer;
    TSharedCounter LaasNoHeader;
    TSharedCounter LocalFails;
    TSharedCounter NoIpInRequest;
    TSharedCounter TrustedRequests;
};


MODULE_WITH_TLS_BASE(geobase, TModuleWithSubModule) {
public:
    TModule(const TModuleParams& mp)
        : TModuleBase(mp)
        , CountersHolders_(mp.Control->SharedStatsManager())
    {
        Config->ForEach(this);

        if (!GeoModule_) {
            ythrow TConfigParseError() << "no geo module configured for geobase";
        }

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

        if (!IpFsm_) { // TODO: get rid of checks after ctor
            ythrow TConfigParseError() << "take_ip_from is required for geobase";
        }

        if (!LaasAnswerHeaderFsm_) {
            ythrow TConfigParseError() << "laas_answer_header is required for geobase";
        }

        TString request = "GET ";
        request += GeoPath_;
        request += " HTTP/1.1\r\n\r\n";
        try {
            TryRethrowError(GeoRequest_.Parse(std::move(request)));
        } catch (...) {
            ythrow TConfigParseError() << "error while building request to geo module in geobase: " << CurrentExceptionMessage();
        }

        UnsafeHeadersFsm_.Reset(new TFsm(*LaasAnswerHeaderFsm_ | TLaasHeaderFsm::Instance()));

        UberFsm_.Reset(new TFsm(
            *IpFsm_ // 0
            | THostFsm::Instance() // 1
            | TXUrlPrefixFsm::Instance() // 2
            | TRestrictedHeadersFsm::Instance() // 3
            | *LaasAnswerHeaderFsm_ // 4
            | TRegionHeaderFsm::Instance() // 5
            | TIpPropertiesHeaderFsm::Instance() // 6
            | TXIsTouristHeaderFsm::Instance() // 7
        ));

        NecessaryHeadersFsm_.Reset(new TFsm(*LaasAnswerHeaderFsm_ // 0
                                     | TRegionHeaderFsm::Instance())); // 1
    }


private:
    START_PARSE {
        PARSE_EVENTS;

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

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

        ON_KEY("geo_host", GeoHost_) {
            return;
        }

        ON_KEY("geo_path", GeoPath_) {
            return;
        }

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

        TString takeIpFrom;
        ON_KEY("take_ip_from", takeIpFrom) {
            IpFsm_.Reset(new TFsm(takeIpFrom, TFsm::TOptions().SetCaseInsensitive(true)));
            return;
        }

        ON_KEY("laas_answer_header", LaasAnswerHeader_) {
            if (!CheckHeaderName(LaasAnswerHeader_)) {
                ythrow TConfigParseError() << "\"laas_answer_header\" is not a valid http header";
            } else if (!CheckRestrictedHeaderName(LaasAnswerHeader_)) {
                ythrow TConfigParseError{} << "\"laas_answer_header\" value " <<  LaasAnswerHeader_.Quote()
                    << " contains one of the restricted headers " << RestrictedHeadersListString() << "\n";
            }
            LaasAnswerHeaderFsm_.Reset(new TFsm(LaasAnswerHeader_, TFsm::TOptions().SetCaseInsensitive(true)));
            return;
        };

        ON_KEY("trusted", MayTrust_) {
            return;
        }

        ON_KEY("processing_time_header", ProcessingTimeHeader_) {
            return;
        }

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

    THolder<TTls> DoInitTls(IWorkerCtl* process) override {
        auto tls = MakeHolder<TTls>(CountersHolders_, process->WorkerId());

        if (!!KillSwitchFile_) {
            tls->KillSwitchChecker = process->SharedFiles()->FileChecker(KillSwitchFile_, TDuration::Seconds(1));
        }
        return tls;
    }

    TError RefineRequest(const TConnDescr& descr, TTls& tls) const noexcept {
        bool succeed = false;
        bool processedOriginalRequest = false;

        if (!tls.KillSwitchFileExists()) {
            const TExtraAccessLogEntry geoLog(descr, "geo");
            if (IsTrusted(descr)) {
                ++tls.TrustedRequests;
                descr.ExtraAccessLog << " trusted";
            } else {
                TParamsCollector collector(descr, GeoRequest_.RequestLine(), GeoHost_, *UberFsm_);

                Y_TRY(TError, error) {
                    Y_PROPAGATE_ERROR(collector.Process(processedOriginalRequest));
                    TRequest request = collector.GeoRequest();

                    TRequester requester(*GeoModule_, descr);
                    TResponse response;
                    Y_PROPAGATE_ERROR(requester.Request(std::move(request), response));

                    succeed = CopyRegionHeaders(response.Headers(), descr.Request->Headers());
                    if (succeed) {
                        descr.ExtraAccessLog << " laas_answered";
                    } else {
                        ++tls.LaasNoHeader;
                    }
                    return {};
                } Y_CATCH {
                    ++tls.LaasNoAnswer;
                    LOG_ERROR(TLOG_ERR, descr, "geo failed: " << GetErrorMessage(error));
                }

                if (!succeed) {
                    descr.ExtraAccessLog << " laas failed";

                    if (!processedOriginalRequest) {
                        if (Y_LIKELY(descr.Request)) {
                            descr.Request->Headers().Delete(*UnsafeHeadersFsm_);
                        }
                    }

                    if (collector.Ok() && collector.IpHeader()) {
                        ++tls.LocalFails;
                    } else {
                        ++tls.NoIpInRequest;
                    }
                }
            }
        } else {
            if (!IsTrusted(descr)) {
                descr.Request->Headers().Delete(*UnsafeHeadersFsm_);
            }
        }

        if (succeed || descr.Request->Headers().FindValues(LaasAnswerHeader_) == descr.Request->Headers().end()) {
            // laas answer header is deleted from descr.Request by preparation fsm's
            // the case where this may be unsafe is when laas_answer_header starts with X-Region-.*,
            // like laas_answer_header = "X-Region-LaaS-Answered", but there is nothing we can do
            // to satisfy everybody
            descr.Request->Headers().Replace(LaasAnswerHeader_, succeed ? "1" : "0");
        }

        return {};
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        TDuration processingTime = TDuration::Max();
        Y_PROPAGATE_ERROR(MeasureProcessingTime([&]() { return RefineRequest(descr, tls); }, &processingTime));

        if (ProcessingTimeHeader_) {
            descr.Request->Headers().Add("X-Yandex-Balancer-LaaS-ProcessingTime", ToString(processingTime.MicroSeconds()));
        }

        return Submodule_->Run(descr);
    }

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

private:
    bool IsTrusted(const TConnDescr& descr) const noexcept {
        return MayTrust_ && HasNecessaryHeaders(descr);
    }

    bool HasNecessaryHeaders(const TConnDescr& descr) const noexcept {
        if (Y_UNLIKELY(!descr.Request)) {
            return false;
        }

        bool hasLaasAnswered{ false };
        bool hasRegion{ false };

        for (const auto& hdr: descr.Request->Headers()) {
            TMatcher matcher{ *NecessaryHeadersFsm_ };

            if (Match(matcher, hdr.first.AsStringBuf()).Final()) {
                switch (*matcher.MatchedRegexps().first) {
                    case 0: // laas_answer_header
                        hasLaasAnswered = true;
                        break;
                    case 1: // X-Region-.*
                        hasRegion = true;
                        break;
                    default:
                        break;
                }
            }

            if (hasLaasAnswered && hasRegion) {
                return true;
            }
        }

        return false;
    }

private:
    THolder<IModule> GeoModule_;
    THolder<TFsm> IpFsm_;

    THolder<TFsm> UberFsm_;
    THolder<TFsm> NecessaryHeadersFsm_;

    TString GeoHost_ = "yabs.yandex.ru";
    TString GeoPath_ = "/region?response_format=header&version=1&service=balancer";
    TString KillSwitchFile_;

    TString LaasAnswerHeader_;
    THolder<TFsm> LaasAnswerHeaderFsm_;
    THolder<TFsm> UnsafeHeadersFsm_;

    TSharedCounters CountersHolders_;

    bool MayTrust_ = false;
    bool ProcessingTimeHeader_ = false;

    TRequest GeoRequest_;
};

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