#include "env.h"

#include <balancer/modules/errordocument/module.h>
#include <balancer/modules/balancer/module.h>
#include <balancer/modules/http/module.h>
#include <balancer/modules/http2/module.h>
#include <balancer/modules/proxy/module.h>
#include <balancer/modules/dummy/module.h>

#include <library/cpp/testing/unittest/env.h>
#include <library/cpp/testing/unittest/registar.h>

#include <utility>


using namespace NBalancerClient;

namespace {
    NJson::TJsonValue BuildSimpleSourceConfig(ui16 status, TString content) {
        NJson::TJsonValue jsonConfig;
        jsonConfig["maxlen"] = 65536;
        jsonConfig["maxreq"] = 65536;
        auto& balancer = jsonConfig["balancer"];
        auto& dynamic = balancer["dynamic"];
        dynamic["max_pessimized_share"] = 0.5;
        auto& section = dynamic["section"];
        NJson::TJsonValue errordoc;
        errordoc["status"] = status;
        errordoc["content"] = std::move(content);
        section["errordocument"] = std::move(errordoc);
        return jsonConfig;
    }

    NSrvKernel::TNodeFactory<NSrvKernel::IModule>& UpdateProxyFactory(NSrvKernel::TNodeFactory<NSrvKernel::IModule>& Factory_) {
        Factory_.AddHandle(NModHttp::Handle());
        Factory_.AddHandle(NModHttp2::Handle());
        Factory_.AddHandle(NModProxy::Handle());
        return Factory_;
    }
}

namespace NBalancerClient::NTesting {
    TBalancerInstanceConfig::TBalancerInstanceConfig(const TMainTaskOptions& options, size_t sdPort) {
        InstanceConfig["enable_reuse_port"] = true;
        InstanceConfig["set_no_file"] = false;
        InstanceConfig["updater_required"] = true;

        if (sdPort) {
            NJson::TJsonValue sdConfig;
            sdConfig["host"] = "localhost";
            sdConfig["client_name"] = "balancer_client_ut";
            sdConfig["port"] = sdPort;
            sdConfig["update_frequency"] = "5ms";

            InstanceConfig["sd"] = sdConfig;
        }

        Globals["workers"] = options.NetworkThreads;
        Globals["maxconn"] = options.MaxConnPerNetworkThread;

        MainOptions.WorkerStartDelay = options.WorkerStartDelay;
        MainOptions.AllowEmptyAdminAddrs = true;
        MainOptions.SetCoroStackSize(200 * 1024);
    }

    TEnv::TEnv()
        : SDServer_(PM_.GetPort())
        , BalancerInstanceConfig_(MainTaskOptions_, SDServer_.GetOptions().Port)
        , MainTask_(MakeHolder<NSrvKernel::NProcessCore::TMainTask>(::ToString(BalancerInstanceConfig_.InstanceConfig),
            BalancerInstanceConfig_.Globals, &Factory_, BalancerInstanceConfig_.MainOptions, nullptr))
    {
        Factory_.AddHandle(NModBalancer::Handle());
        Factory_.AddHandle(NModErrorDocument::Handle());
        Factory_.AddHandle(NModHttp::Handle());
        Factory_.AddHandle(NModProxy::Handle());
        Factory_.AddHandle(NModDummy::Handle());
    }

    void TEnv::Start() {
        Started_ = true;
        SDServer_.Start();

        MainTaskPool_.Start(1, 0);
        UNIT_ASSERT(MainTaskPool_.AddFunc([this]() {
            try {
                MainTask_->Execute();
            } catch (...) {
                UNIT_ASSERT_C(false, CurrentExceptionMessage());
                throw;
            }
        }));
    }

    TEnv::~TEnv() {
        NameToSource_.clear();
        if (Started_) {
            MainTask_->StopMaster({
                .CoolDown = TDuration{},
                .Timeout = TDuration{},
                .CloseTimeout = TDuration::Seconds(10),
            });
            MainTaskPool_.Stop();
            SDServer_.Stop();
        }
    }

    THolder<TBalancerClient> TEnv::CreateSource(const NJson::TJsonValue& config) {
        with_lock (Lock_) {
            return MakeHolder<TBalancerClient>(MainTask_, Factory_, "http", config);
        }
    }

    void TEnv::DeleteSource(const TString& name) {
        with_lock (Lock_) {
            Y_ASSERT(NameToSource_[name]);
            NameToSource_.erase(name);
        }
    }

    void TEnv::CreateSimpleSource(const TString& name, ui16 status, const TString& content) {
        with_lock (Lock_) {
            Y_ASSERT(!NameToSource_[name]);
            NameToSource_[name] = MakeHolder<TBalancerClient>(MainTask_, Factory_, "http", BuildSimpleSourceConfig(status, content));
        }
    }

    THolder<TRequestContext> TEnv::SendRequest(const TString& source, TRequest request, TString data) {
        Y_ASSERT(NameToSource_[source]);
        return NameToSource_[source]->SendRequest(std::move(request), std::move(data), MakeHolder<NSrvKernel::TAttemptsHolderBase>(1, 0), nullptr);
    }

    void TEnv::WaitStart() {
        const TInstant startDeadline = Now() + TDuration::Seconds(10);
        size_t workersReady = 0;

        while (Now() < startDeadline) {
            workersReady = AtomicGet(MainTask_->LiveWorkersCounter_);
            if (workersReady == MainTaskOptions_.NetworkThreads) {
                break;
            }

            Sleep(TDuration::MilliSeconds(10));
        }

        if (workersReady != MainTaskOptions_.NetworkThreads) {
            ythrow yexception() << "start only " << workersReady << " workers";
        }
    }
}


// NBalancerClient::NTesting::TBalancerProxyInstanceConfig
namespace NBalancerClient::NTesting {
    TBalancerProxyInstanceConfig::TBalancerProxyInstanceConfig(ui16 proxyPort, ui16 serverPort, bool http2Backend) {
        InstanceConfig["enable_reuse_port"] = true;
        InstanceConfig["set_no_file"] = false;
        InstanceConfig["updater_required"] = true;

        NJson::TJsonValue addrs;
        addrs["ip"] = "localhost";
        addrs["port"] = proxyPort;

        InstanceConfig["addrs"].AppendValue(addrs);

        InstanceConfig["http2"]["http"]["maxlen"] = 8096;
        InstanceConfig["http2"]["http"]["maxreq"] = 8096;

        InstanceConfig["http2"]["http"]["proxy"]["backend_timeout"] = "5s";

        InstanceConfig["http2"]["http"]["proxy"]["host"] = "localhost";
        InstanceConfig["http2"]["allow_http2_without_ssl"] = true;
        InstanceConfig["http2"]["allow_sending_trailers"] = true;
        InstanceConfig["http2"]["http"]["proxy"]["port"] = serverPort;

        InstanceConfig["http2"]["http"]["proxy"]["http2_backend"] = http2Backend;

        Globals["workers"] = MainTaskOptions.NetworkThreads;
        Globals["maxconn"] = MainTaskOptions.MaxConnPerNetworkThread;

        MainOptions.WorkerStartDelay = MainTaskOptions.WorkerStartDelay;
        MainOptions.AllowEmptyAdminAddrs = true;
        MainOptions.SetCoroStackSize(200 * 1024);
    }

    TProxy::TProxy(TBalancerProxyInstanceConfig&& config)
        : BalancerProxyInstanceConfig_(config)
        , MainTask_(MakeHolder<NSrvKernel::NProcessCore::TMainTask>(::ToString(BalancerProxyInstanceConfig_.InstanceConfig),
            BalancerProxyInstanceConfig_.Globals, &UpdateProxyFactory(Factory_), BalancerProxyInstanceConfig_.MainOptions, nullptr))
    {
    }


    void TProxy::Start() {
        MainTaskPool_.Start(1, 0);
        Y_ENSURE(MainTaskPool_.AddFunc([this]() {
            try {
                MainTask_->Execute();
            } catch (...) {
                throw;
            }
        }));
        WaitStart();
    }

    void TProxy::Stop() {
        MainTask_->StopMaster({
            .CoolDown = TDuration{},
            .Timeout = TDuration{},
            .CloseTimeout = TDuration::Seconds(10),
            });
        MainTaskPool_.Stop();
    }


    TProxy::~TProxy() {
        Stop();
    }

    void TProxy::WaitStart() {
            const TInstant startDeadline = Now() + TDuration::Seconds(10);
            size_t workersReady = 0;

            while (Now() < startDeadline) {
                workersReady = AtomicGet(MainTask_->LiveWorkersCounter_);
                if (workersReady == MainTaskOptions_.NetworkThreads) {
                    break;
                }

                Sleep(TDuration::MilliSeconds(10));
            }

            if (workersReady != MainTaskOptions_.NetworkThreads) {
                ythrow yexception() << "start only " << workersReady << " workers";
            }
    }
}
