#include "module.h"
#include "transfer.h"
#include "config.h"
#include "tls.h"
#include "proxy.h"

#include <balancer/kernel/http2/client/http2_backend.h>
#include <balancer/kernel/http2/client/http2_request.h>
#include <balancer/kernel/client_request/backend.h>
#include <balancer/kernel/client_request/backend_config.h>
#include <balancer/kernel/client_request/face.h>
#include <balancer/kernel/connection_manager/connection_manager.h>
#include <balancer/kernel/connection_manager/connection_pool.h>
#include <balancer/kernel/connection_manager_helpers/helpers.h>
#include <balancer/kernel/custom_io/countio.h>
#include <balancer/kernel/dns/async_resolver.h>
#include <balancer/kernel/dns/resolver.h>
#include <balancer/kernel/http/parser/httpdecoder.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/http/parser/status_codes/status_codes.h>
#include <balancer/kernel/http2/server/common/http2_common.h>
#include <balancer/kernel/log/errorlog.h>
#include <balancer/kernel/module/module.h>
#include <balancer/kernel/net/sockops.h>
#include <balancer/kernel/process/main_config.h>
#include <balancer/kernel/requester/requester.h>
#include <balancer/kernel/srcrwr/srcrwr_addrs.h>

#include <library/cpp/containers/stack_vector/stack_vec.h>

#include <util/generic/scope.h>
#include <util/random/shuffle.h>
#include <util/string/builder.h>

using namespace NConfig;
using namespace NSrvKernel;
using namespace NModProxy;

TAtomic ProxyIdCounter_ = 0;

Y_TLS(proxy), public NModProxy::TTls {
  public:
    TTls(NSrvKernel::TSharedStatsManager& statsManager,
         NConnectionManager::TConnectionPool* connectionPool,
         IWorkerCtl* process,
         size_t keepAliveCount,
         TMaybe<TDuration> keepAliveTimeout
    )
        : NModProxy::TTls{statsManager}
        , ConnectionPool_{connectionPool}
        , Process_{process}
        , KeepAliveCount_{keepAliveCount}
        , KeepAliveTimeout_{keepAliveTimeout.GetOrElse(TDuration::Zero())}
        , UnusedKeepalives_(statsManager.MakeGauge("proxy-unused_keepalives").AllowDuplicate().Build())
    {
        if (KeepAliveCount_) {
            if (ConnectionPool_) {
                WorkerConnectionPool_ = connectionPool->GetPoolForWorker(process);
            } else if (KeepAliveTimeout_ != TDuration::Zero()) {
                KeepAliveCleanerTask_ = NSrvKernel::TCoroutine{NSrvKernel::ECoroType::Service,
                                                              "proxy_keepalive_cleaner_cont",
                                                              &process->Executor(),
                                                              [this] {
                                                                  CleanKeepalives();
                                                              }};
            }
        }
    }

    void Dispose() {
        WorkerConnectionPool_ = nullptr;
        if (ConnectionPool_) {
            ConnectionPool_->ReleasePoolForWorker(Process_);
        }
        UnusedKeepalives_.Sub(KeepAliveConns_.size());
        KeepAliveConns_.clear();
    }

    THolder<TKeepAliveData> GetKeepAliveConnection(const TConnDescr& descr, const THostInfo& hostInfo) noexcept final {
        THolder<TKeepAliveData> keepAlive;

        if (!CanStoreKeepaliveConnection(descr, hostInfo)) {
            return {};
        }

        if (WorkerConnectionPool_) {
            keepAlive = WorkerConnectionPool_->TryPullConnection();
        } else if (!KeepAliveConns_.empty()) {
            keepAlive = std::move(KeepAliveConns_.front());
            KeepAliveConns_.pop_front();
        }

        if (keepAlive) {
            UnusedKeepalives_.Dec();
        }

        return keepAlive;
    }

    bool CanStoreKeepaliveConnection(const TConnDescr& descr, const THostInfo& hostInfo) const noexcept final {
        return descr.Request && !descr.Properties->Parent.SkipKeepalive && !hostInfo.IsSrcRwr;
    }

    void StoreKeepaliveConnection(THolder<TKeepAliveData> keepAlive) noexcept final {
        if (WorkerConnectionPool_) {
            if (WorkerConnectionPool_->AddConnection(keepAlive)) {
                UnusedKeepalives_.Inc();
            }
        } else {
            if (KeepAliveConns_.size() < KeepAliveCount_) {
                if (KeepAliveTimeout_ != TDuration::Zero()) {
                    keepAlive->SetKeepAliveReleaseDeadline(KeepAliveTimeout_.ToDeadLine());
                }
                KeepAliveConns_.push_front(std::move(keepAlive));
                UnusedKeepalives_.Inc();

                if (KeepAliveConns_.size() == 1) {
                    KeepAliveCleanerCondVar_.notify();
                }
            }
        }
    }

  private:
    void CleanKeepalives() {
        auto cont = RunningCont();
        auto executor = cont->Executor();
        do {
            while (!KeepAliveConns_.empty()) {
                auto connection = std::move(KeepAliveConns_.back());
                KeepAliveConns_.pop_back();

                Y_ASSERT(connection);
                if (connection->KeepAliveReleaseDeadline() <= Now() ||
                    !IsNotSocketClosedByOtherSide(connection->Sock())) {
                    connection.Destroy();
                    UnusedKeepalives_.Dec();
                } else {
                    KeepAliveConns_.push_back(std::move(connection));
                    break;
                }

                if (cont->Cancelled()) {
                    break;
                }
            }

            if (KeepAliveConns_.empty()) {
                const int status = KeepAliveCleanerCondVar_.wait_for(executor, KeepAliveTimeout_);
                if (status == ECANCELED) {
                    break;
                }
            } else {
                const int status = KeepAliveCleanerCondVar_.wait_until(
                        executor, KeepAliveConns_.back()->KeepAliveReleaseDeadline());
                if (status == ECANCELED) {
                    break;
                }
            }
        } while (!cont->Cancelled());
    }

  private:
    NConnectionManager::TConnectionPool* ConnectionPool_;
    IWorkerCtl* Process_;
    size_t KeepAliveCount_ = 0;
    TDuration KeepAliveTimeout_;
    std::deque<THolder<NModProxy::TKeepAliveData>> KeepAliveConns_;
    NSrvKernel::TSharedCounter UnusedKeepalives_;
    NConnectionManager::TWorkerConnectionPool* WorkerConnectionPool_ = nullptr;
    NSrvKernel::TCoroSingleCondVar KeepAliveCleanerCondVar_;
    NSrvKernel::TCoroutine KeepAliveCleanerTask_;
};

MODULE_WITH_TLS(proxy) {
public:
    explicit TModule(const TModuleParams& mp)
        : TModuleBase(mp)
        , UnusedKeepalives_(mp.Control->SharedStatsManager().MakeGauge("proxy-unused_keepalives").AllowDuplicate().Build())
    {
        ProxyConfig_.Name = NAME;
        ProxyConfig_.StatusCodeReactions.AddClassReaction(5, NSrvKernel::TStatusCodeReaction::Bad()); // default fail_on_5xx = true

        ProxyConfig_.BackendConfig = ParseProtoConfig<TBackendConfig>(
            [&](const TString&, NConfig::IConfig::IValue*) {
            }
        );

        if (ProxyConfig_.BackendConfig.has_https_settings()) {
            InitSsl();
        }

        Config->ForEach(this);

        if (!ProxyConfig_.BackendConfig.host()) {
            ythrow TConfigParseError() << "no host configured for proxy";
        }

        if (!ProxyConfig_.BackendConfig.need_resolve() && !ProxyConfig_.BackendConfig.cached_ip()) {
            ythrow TConfigParseError() << "need_resolve requires cached_ip to be set";
        }

        if (!ProxyConfig_.BackendConfig.cached_ip() && mp.GetResolvedHosts()) {
            auto iter = mp.GetResolvedHosts()->find(ProxyConfig_.BackendConfig.host());
            if (iter != mp.GetResolvedHosts()->end() && !iter->second.empty()) {
                ProxyConfig_.BackendConfig.set_cached_ip(iter->second);
            }
        }

        if (ProxyConfig_.BackendConfig.cached_ip()) {
            TError error = ProxyConfig_.BackendConfig.FillCachedAddr();
            if (!ProxyConfig_.BackendConfig.GetCachedAddr() || ProxyConfig_.BackendConfig.GetCachedAddr()->Addresses.empty()) {
                TString explanation = error ? GetErrorMessage(error) : "unknown error";
                ythrow TConfigParseError() << "can't parse cached_ip " << ProxyConfig_.BackendConfig.cached_ip().Quote() << ", " << explanation;
            }
        }

        if (ProxyConfig_.BackendConfig.has_connect_retry_timeout() && ProxyConfig_.BackendConfig.connect_retry_timeout() <= ProxyConfig_.BackendConfig.connect_retry_delay()) {
            Y_WARN_ONCE("\"connect_retry_timeout\" = " << ProxyConfig_.BackendConfig.connect_retry_timeout() << " disables \"connect_retry_delay\" = " << ProxyConfig_.BackendConfig.connect_retry_delay());
        }

        if (ProxyConfig_.BackendConfig.backend_timeout() == TDuration::Max() && HasHttpParent()) {
            Y_WARN_ONCE("\"backend_timeout\" is not specified, setting backend_timeout = 5s");
            ProxyConfig_.BackendConfig.set_backend_timeout(TDuration::Seconds(5)); // see https://st.yandex-team.ru/BALANCER-1296
        }

        if (ProxyConfig_.BackendConfig.connect_timeout() == TDuration::Max()) {
            Y_WARN_ONCE("\"connect_timeout\" is not specified");
        }

        if (ProxyConfig_.BackendConfig.resolve_timeout() == TDuration::Max()) {
            Y_WARN_ONCE("\"resolve_timeout\" is not specified");
        }

        if (ProxyConfig_.BackendConfig.GetSockBufSize().Rcv) {
            // For test purpose only
            Y_WARN_ONCE("\"_sock_inbufsize\" may not work as you expect");
        }

        if (ProxyConfig_.BackendConfig.GetSockBufSize().Snd) {
            // For test purpose only
            Y_WARN_ONCE("\"_sock_outbufsize\" may not work as you expect");
        }

        if (ProxyConfig_.KeepAliveTimeout && ProxyConfig_.KeepAliveCount <= 0) {
            ythrow TConfigParseError() << "\"keepalive_timeout\" requires \"keepalive_count\" to be positive number";
        }

        if (ProxyConfig_.KeepAliveCount > 0 && !ProxyConfig_.KeepAliveTimeout.Defined()) {
            ProxyConfig_.KeepAliveTimeout = TDuration::Minutes(5);
        }

        if (ProxyConfig_.AllowConnectionUpgradeWithoutConnectionHeader && !ProxyConfig_.AllowConnectionUpgrade) {
            ythrow TConfigParseError{} << "allow_connection_upgrade_without_connection_header used without allow_connection_upgrade";
        }

        TryRethrowError(ValidateKeepalive(ProxyConfig_.BackendConfig.GetTcpKeepalive()));
        TryRethrowError(ValidateSockBufSize(ProxyConfig_.BackendConfig.GetSockBufSize()));

        if (ProxyConfig_.KeepAliveCount > 0 && Control->ConnectionManager()) {
            ProxyId_ = AtomicGetAndIncrement(ProxyIdCounter_);
            ConnectionPool_ = Control->ConnectionManager()->AddConnectionPool(
                    ProxyId_, ProxyConfig_.KeepAliveCount, PM_EDGE_TRIGGERED,
                    ProxyConfig_.KeepAliveTimeout.GetOrElse(TDuration::Zero()),
                ProxyConfig_.BackendConfig.connect_timeout() == TDuration::Max() ? TDuration::MilliSeconds(50) : ProxyConfig_.BackendConfig.connect_timeout());
            ConnectionPool_->SetOnExpire([counter = MakeAtomicShared<TSharedCounter>(UnusedKeepalives_, 0u)](){
                counter->Dec();
            });
        }

        Impl_ = MakeHolder<TModuleProxy>(ProxyConfig_);
    }

    ~TModule() {
        if (ConnectionPool_) {
            Control->ConnectionManager()->RemoveConnectionPool(ProxyId_);
        }
    }

private:
    START_PARSE {
        if (key == "host") {
            return;
        }

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

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

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

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

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

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

        if (IsIn({"use_only_ipv4", "use_only_ipv6"}, key)) {
            return;
        }

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

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

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

        if (TStringBuf(key).StartsWith("tcp_keep_")) {
            return;
        }

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

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

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

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

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

        ON_KEY("switched_backend_timeout", ProxyConfig_.SwitchedBackendTimeout) {
            return;
        }

        ON_KEY("backend_read_timeout", ProxyConfig_.BackendReadTimeout) {
            return;
        }

        ON_KEY("backend_write_timeout", ProxyConfig_.BackendWriteTimeout) {
            return;
        }

        ON_KEY("client_read_timeout", ProxyConfig_.ClientReadTimeout) {
            return;
        }

        ON_KEY("client_write_timeout", ProxyConfig_.ClientWriteTimeout) {
            return;
        }

        if (key == "fail_on_empty_reply") {
            Y_WARN_ONCE("\"fail_on_empty_reply\" will be erased soon");
            return;
        }

        bool failOn5xx = false ;
        ON_KEY("fail_on_5xx", failOn5xx) {
            ProxyConfig_.StatusCodeReactions.AddClassReaction(5, failOn5xx ? TStatusCodeReaction::Bad() : TStatusCodeReaction::Good());
            return;
        }

        if (key == "status_code_blacklist") {
            TStatusCodeReactionParser parser(ProxyConfig_.StatusCodeReactions, TStatusCodeReaction::Bad());
            value->AsSubConfig()->ForEach(&parser);
            return;
        }

        if (key == "status_code_blacklist_exceptions") {
            TStatusCodeReactionParser parser(ProxyConfig_.StatusCodeReactions, TStatusCodeReaction::Good());
            value->AsSubConfig()->ForEach(&parser);
            return;
        }

        ON_KEY("buffering", ProxyConfig_.Buffering) {
            return;
        }

        ON_KEY("keepalive_timeout", ProxyConfig_.KeepAliveTimeout) {
            return;
        }

        if (key == "keepalive_count") {
            if (value->AsString() == "unlimited") {
                ProxyConfig_.KeepAliveCount = Max<size_t>();
                return;
            }
            if (NConfig::ParseFromString(value->AsString(), ProxyConfig_.KeepAliveCount)) {
                return;
            }
        }

        ON_KEY("keepalive_check_for_unexpected_data", ProxyConfig_.KeepAliveCheckForUnexpectedData) {
            return;
        }

        // TODO(velavokr): BALANCER-2342 remove completely after new L7 configs a released
        if (key == "wait_unsent_to_client") {
            return;
        }

        // TODO(velavokr): BALANCER-2318 remove completely after new L7 configs a released
        if (key == "count_body_length") {
            return;
        }

        ON_KEY("allow_connection_upgrade", ProxyConfig_.AllowConnectionUpgrade) {
            return;
        }

        ON_KEY("allow_connection_upgrade_without_connection_header", ProxyConfig_.AllowConnectionUpgradeWithoutConnectionHeader) {
            return;
        }

        ON_KEY("watch_client_close", ProxyConfig_.WatchClientClose) {
            return;
        }

        ON_KEY("keep_all_headers", ProxyConfig_.KeepAllHeaders) {
            return;
        }

        ON_KEY("http2_backend", ProxyConfig_.Http2Backend) {
            return;
        }

        if (key == "ignore_backend_write_error_file") {
            return;
        }
        if (key == "check_backend_write_error") {
            return;
        }
        if (key == "ignore_http_client_read_error_file") {
            return;
        }
        if (key == "check_http_client_read_error") {
            return;
        }

        if (key == "http_backend") {
            return;
        }
    } END_PARSE

private:
    THolder<TTls> DoInitTls(IWorkerCtl* process) noexcept override {
        auto tls = MakeHolder<TTls>(Control->SharedStatsManager(), ConnectionPool_, process, ProxyConfig_.KeepAliveCount, ProxyConfig_.KeepAliveTimeout);

        for (const auto& protocol : ProxyConfig_.BackendConfig.protocols()) {
            tls->State.AddImpl(Control->GetProtocolFactory(), protocol, ProxyConfig_.BackendConfig);
        }
        // Adding HTTP2 protocol support
        TMaybe<TBackendProtocolFactory> backendProtocolFactory;
        if (!Control->GetProtocolFactory().HasProtocolImpl<NBalancerClient::THttp2Backend>())
        {
            backendProtocolFactory.ConstructInPlace();
            backendProtocolFactory->RegisterProtocolImpl<NBalancerClient::THttp2Backend>();
        }

        tls->State.AddImpl(backendProtocolFactory ? *backendProtocolFactory : Control->GetProtocolFactory(), NBalancerClient::THttp2Backend::ProtocolName(), ProxyConfig_.BackendConfig);
        return tls;
    }

    TError DoRun(const TConnDescr& descr, TTls& tls) const noexcept override {
        return Impl_->Run(descr, tls);
    }

    void DoDispose(IWorkerCtl* /*process*/, TTls& tls) noexcept override {
        tls.Dispose();
    }

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

    bool HasHttpParent() const noexcept {
        return CheckParents([&](TStringBuf name) {
            return "http" == name;
        });
    }

    TBackendCheckResult CheckBackends(IWorkerCtl& proc, bool runtimeCheck, TTls&) const noexcept override {
        if (runtimeCheck) {
            return TBackendCheckResult{TBackendCheckResult::EStatus::Skipped};
        }
        TAsyncRequester requester(*this, nullptr, proc);
        requester.Props().SkipKeepalive = true;
        TNullStream in;
        TNullStream out;
        TBackendCheckResult result{TBackendCheckResult::EStatus::Success};
        if (TError err = requester.Requester().TcpRequest(in, out)) {
            result.Errors.emplace_back(std::move(err));
            result.Status = TBackendCheckResult::EStatus::Failed;
        }
        return result;
    }

public:
    void PrintInfo(NJson::TJsonWriter& out) const noexcept override {
        out.WriteKey("proxy");
        out.OpenMap();
        out.Write("host", ProxyConfig_.BackendConfig.host());
        out.Write("port", ProxyConfig_.BackendConfig.port());
        out.CloseMap();
    }

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

private:
    TProxyConfig ProxyConfig_;
    THolder<TModuleProxy> Impl_;
    NConnectionManager::TConnectionPool* ConnectionPool_ = nullptr;

    //TODO: make these counters hosted by overlying report modules instead of global
    TSharedCounter UnusedKeepalives_;

    size_t ProxyId_ = 0;
    TMutex Mutex_;
};

IModuleHandle* NModProxy::Handle() {
    return Nproxy::TModule::Handle();
}
