#include "server.h"

#include <balancer/kernel/process/master_main.h>

#include <balancer/server/config.h>
#include <balancer/server/modules/common.h>
#include <balancer/server/ut/util/common.h>

#include <library/cpp/coroutine/engine/helper.h>
#include <library/cpp/neh/neh.h>
#include <library/cpp/neh/http_common.h>

#include <library/cpp/testing/unittest/tests_data.h>

#include <util/thread/pool.h>

using namespace NBalancerServer;
using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TestServersOnSameWorkers) {
    struct TEnv3 {
        struct TServerWithOptions {
            TServerWithOptions(TPortManager& portManager, TOptions options,
                               const TServerCallback& callback,
                               TAtomicSharedPtr<NSrvKernel::NProcessCore::TMainTask>& mainTask, bool useTwoPorts)
                : Port(portManager.GetPort())
                , SslPort(portManager.GetPort())
                , Options(std::move(options))
            {
                if (useTwoPorts) {
                    Options.SslPort = SslPort;
                    Options.Port = Port;
                } else if (Options.SslConfigPath) {
                    Options.SslPort = SslPort;
                } else {
                    Options.Port = Port;
                }

                Server = MakeHolder<TServer>(callback, mainTask, Options);
            }

            const ui32 Port = 0;
            const ui32 SslPort = 0;
            TOptions Options;
            THolder<TServer> Server;
        };

        TPortManager PortManager;

        THolder<TBalancerInstanceConfig> InstanceConfig;
        NSrvKernel::TNodeFactory<NSrvKernel::IModule> ModulesFactory;
        TAtomicSharedPtr<NSrvKernel::NProcessCore::TMainTask> MainTask;
        TVector<TServerWithOptions> Servers;
        TThreadPool ServerPool;

        void Start(const TServerCallback& callback, const TOptions& options, bool useTwoPorts = false) {
            InstanceConfig = MakeHolder<TBalancerInstanceConfig>(options);

            AddCommonServerModules(ModulesFactory);

            MainTask = MakeHolder<NSrvKernel::NProcessCore::TMainTask>(InstanceConfig->LuaInstanceConfig, InstanceConfig->Globals,
                                                                       &ModulesFactory, InstanceConfig->MainOptions, nullptr);

            for (size_t i = 0; i < 3; ++i) {
                Servers.emplace_back(PortManager, options, callback, MainTask, useTwoPorts);
            }

            ServerPool.Start(1, 0);

            for (const auto& s : Servers) {
                UNIT_ASSERT(!NCoro::TryConnect("localhost", s.Port));
                UNIT_ASSERT(!NCoro::TryConnect("localhost", s.SslPort));
                s.Server->Run();
            }

            UNIT_ASSERT(ServerPool.AddFunc([this]() {
                try {
                    this->MainTask->Execute();
                } catch (...) {
                    UNIT_ASSERT_C(false, CurrentExceptionMessage());
                    throw;
                }
            }));

            for (const auto& s : Servers) {
                if (s.Options.Port) {
                    UNIT_ASSERT(NCoro::WaitUntilConnectable("localhost", s.Options.Port, TDuration::Seconds(10)));
                }

                if (s.Options.SslPort) {
                    UNIT_ASSERT(NCoro::WaitUntilConnectable("localhost", s.Options.SslPort, TDuration::Seconds(10)));
                }
            }
        }

        void Stop() {
            MainTask->StopMaster({
                .CoolDown = TDuration{},
                .Timeout = TDuration{},
                .CloseTimeout = TDuration::Seconds(10),
            });
            ServerPool.Stop();
            for (const auto& s : Servers) {
                s.Server->Wait();
                UNIT_ASSERT(!NCoro::TryConnect("localhost", s.Port));
                UNIT_ASSERT(!NCoro::TryConnect("localhost", s.SslPort));
            }
        }

        ~TEnv3() {
            Stop();
        }
    };

    Y_UNIT_TEST(TestOk3Servers) {
        for (size_t threads : {0, 1, 5}) {
            NTesting::RunConcurrent<TEnv3>(
                NTesting::GetOkCallback(),
                [](TEnv3& env) {
                    for (const auto& i : env.Servers) {
                        auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << i.Port << "/yandsearch?cgi=1")->Wait();
                        UNIT_ASSERT_C(!res->IsError(), res->GetErrorText());
                        UNIT_ASSERT_VALUES_EQUAL(res->Data, "OK");
                    }
                },
                TOptions().SetThreads(threads)
            );
        }
    }
}

