#include <balancer/client/ut/util/env.h>
#include <balancer/client/ut/util/service.h>

#include <balancer/kernel/http/parser/request_builder.h>
#include <balancer/modules/balancer/module.h>

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

using namespace NSrvKernel;
using namespace NBalancerClient;
using namespace NBalancerClient::NTesting;

using NThreading::TFuture;
using NThreading::TPromise;

namespace {

    const TString CertPath = ArcadiaSourceRoot() + "/balancer/test/plugin/certs/data/default.crt";
    const TString PrivPath = ArcadiaSourceRoot() + "/balancer/test/plugin/certs/data/default.key";
    const TString CaPath = ArcadiaSourceRoot() + "/balancer/test/plugin/certs/data/root_ca.crt";

    void MakeSslConfig(IOutputStream& out) {
        out << "contexts = {\n";
        out << "  default = {\n";
        out << "     cert = " << CertPath.Quote() << ";\n";
        out << "     priv = " << PrivPath.Quote() << ";\n";
        out << "     ciphers = \"kEECDH:kRSA+AES128:kRSA:+3DES:RC4:!aNULL:!eNULL:!MD5:!EXPORT:!LOW:!SEED:!CAMELLIA:!IDEA:!PSK:!SRP:!SSLv2\";\n";
        out << "  };\n";
        out << "};\n";
    }

}  // namespace


Y_UNIT_TEST_SUITE(TestClient) {
    Y_UNIT_TEST(Trivial) {
        TEnv env;
        env.Start();
    }

    struct TSimpleSourceDescr {
        TString Name;
        ui16 Status;
        TString Content;
    };

    void MakeReqs(TEnv& env, const TVector<TSimpleSourceDescr>& descrs, size_t count) {
        TVector<THolder<TRequestContext>> contexts(count * descrs.size());
        TVector<TVector<TFuture<TBalancerClientResponseContext>>> futuresOfSource(descrs.size());

        for (size_t i = 0u; i < count; ++i) {
            for (size_t s = 0u; s < descrs.size(); ++s) {
                contexts.push_back(env.SendRequest(descrs[s].Name, TRequest(), ""));
                futuresOfSource[s].push_back(contexts.back()->GetFuture());
            }
        }

        for (size_t s = 0u; s < descrs.size(); ++s) {
            WaitAll(futuresOfSource[s]).Wait();
            for (auto& i : futuresOfSource[s]) {
                TBalancerClientResponseContext respContext = i.ExtractValueSync();
                TBalancerClientResponse resp{};
                auto error = respContext.GetResponse().AssignTo(resp);
                UNIT_ASSERT_C(!error, error->what());
                UNIT_ASSERT_VALUES_EQUAL(resp.HttpResponse.ResponseLine().StatusCode, descrs[s].Status);
                UNIT_ASSERT_VALUES_EQUAL(resp.Data, descrs[s].Content);
            }
        }
    }

    Y_UNIT_TEST(CreateSourceBeforeStart) {
        TEnv env;

        TSimpleSourceDescr descr{"succ", 200, ":)"};
        env.CreateSimpleSource(descr.Name, descr.Status, descr.Content);

        env.Start();

        MakeReqs(env, {descr}, 100);
    }

    Y_UNIT_TEST(CreateSourceAfterStart) {
        TEnv env;

        env.Start();

        TSimpleSourceDescr descr{"succ", 200, ":)"};
        env.CreateSimpleSource(descr.Name, descr.Status, descr.Content);

        MakeReqs(env, {descr}, 100);
    }

    Y_UNIT_TEST(CreateSourcesBeforeAndAfterStart) {
        TEnv env;

        TSimpleSourceDescr descr1{"succ", 200, ":)"};
        env.CreateSimpleSource(descr1.Name, descr1.Status, descr1.Content);

        env.Start();

        TSimpleSourceDescr descr2{"fail", 500, ":("};
        env.CreateSimpleSource(descr2.Name, descr2.Status, descr2.Content);

        MakeReqs(env, {descr1, descr2}, 100);
    }

    Y_UNIT_TEST(CreateAndDeleteSourceBeforeStart) {
        TEnv env;

        TSimpleSourceDescr descr{"succ", 200, ":)"};
        env.CreateSimpleSource(descr.Name, descr.Status, descr.Content);
        env.DeleteSource(descr.Name);

        env.Start();
    }

    Y_UNIT_TEST(CreateBeforeAndDeleteAfterStart) {
        TEnv env;

        TSimpleSourceDescr descr{"succ", 200, ":)"};
        env.CreateSimpleSource(descr.Name, descr.Status, descr.Content);

        env.Start();

        MakeReqs(env, {descr}, 100);

        env.DeleteSource(descr.Name);
    }


    NSrvKernel::TError EchoReply(NBalancerServer::THttpRequestEnv& env) {
        auto reply = env.GetReplyTransport();

        NSrvKernel::TResponse head(200, "Ok");
        head.Props().ContentLength = env.Body().size();
        reply->SendHead(std::move(head));
        reply->SendData(env.Body());
        reply->SendEof();

        return NSrvKernel::TError();
    }

    struct TConcurrentTest {
        TThreadPool Pool;

        TMutex Lock;
        TVector<TAtomicSharedPtr<TBalancerClient>> Clients{3};

        TAtomic ReqCounter = 0;
        TAtomic Stopped = 0;
    public:
        void Stop() {
            Pool.Stop();
        }

        void ScheduleRequests(TEnv& env, const NJson::TJsonValue& config, bool useEncryption = false, bool useStreaming = false) {
            const size_t nThreads = 10;

            Pool.Start(nThreads);

            for (size_t i = 0; i < nThreads; ++i) {
                Y_VERIFY(Pool.AddFunc([this, &env, &config, useEncryption, useStreaming]() {
                    try {
                        THolder<TBalancerClient> client;
                        while (!AtomicGet(Stopped)) {
                            TAtomicSharedPtr<TBalancerClient> client;

                            with_lock (Lock) {
                                auto r = AtomicIncrement(ReqCounter) % Clients.size();
                                auto& c = Clients[r];
                                if (!c || r % 100 == 0) {
                                   c = env.CreateSource(config);
                                }
                                client = c;
                            }

                            size_t dataSize = RandomNumber<size_t>(1000000);
                            TString data = NUnitTest::RandomString(dataSize, dataSize);

                            NSrvKernel::TRequest request = BuildRequest().Method(EMethod::POST).Version11().Path("/yandsearch?").Cgi("text=1");
                            request.Props().ContentLength = data.size();

                            TAtomicSharedPtr<TU2WChannel<TString>> inputChannel = MakeAtomicShared<TU2WChannel<TString>>(100ul);

                            auto context = client->SendRequest(std::move(request), useStreaming ? TString() : data,
                                MakeHolder<NSrvKernel::TAttemptsHolderBase>(1, 0, useEncryption),
                                    useStreaming ? inputChannel : nullptr);

                            if (useStreaming) {
                                size_t sent = 0;
                                while (sent != data.size()) {
                                    size_t length = 1 + (data.size() - sent - 1 ? RandomNumber<size_t>(data.size() - sent - 1) : 0);
                                    TString chunk(data.begin() + sent, data.begin() + sent + length);
                                    EChannelStatus status = inputChannel->Send(std::move(chunk), TInstant::Max());
                                    UNIT_ASSERT(status == EChannelStatus::Success);
                                    sent += length;
                                }
                            }

                            TBalancerClientResponse response;
                            if (auto error = context->GetFuture().ExtractValueSync().GetResponse().AssignTo(response)) {
                                UNIT_ASSERT_C(false, error->what());
                            } else {
                                UNIT_ASSERT_VALUES_EQUAL(response.HttpResponse.ResponseLine().StatusCode, 200);
                                UNIT_ASSERT_VALUES_EQUAL(response.Data, data);
                            }
                        }
                    } catch (...) {
                        UNIT_ASSERT_C(false, CurrentExceptionMessage());
                    }
                }));
            }
        }
    };

    Y_UNIT_TEST(ProxyClient) {
        TEnv env;
        env.Start();

        TAtomic requestCount = 0;

        TService echoService([&requestCount](NBalancerServer::THttpRequestEnv& env) {
                AtomicIncrement(requestCount);
                return EchoReply(env);
            },
            {}
        );

        NJson::TJsonValue config;
        config["maxlen"] = 65536;
        config["maxreq"] = 65536;
        auto& dynamic = config["balancer"]["dynamic"];
        dynamic["max_pessimized_share"] = 0.5;

        NJson::TJsonValue proxy;
        proxy["cached_ip"] = "127.0.0.1";
        proxy["host"] = "localhost";
        proxy["port"] = echoService.Port;
        proxy["backend_timeout"] = "1s";

        NJson::TJsonValue item;
        item["proxy"] = proxy;
        dynamic.AppendValue(item);

        TConcurrentTest ct;
        ct.ScheduleRequests(env, config);

        while (AtomicGet(ct.ReqCounter) < 1000) {
            Sleep(TDuration::MilliSeconds(10));
        }

        AtomicSet(ct.Stopped, 1);
        ct.Stop();

        UNIT_ASSERT_VALUES_EQUAL(requestCount, ct.ReqCounter);    }

    Y_UNIT_TEST(ProxyClientWithStreaming) {
        TEnv env;
        env.Start();

        TAtomic requestCount = 0;

        TService echoService([&requestCount](NBalancerServer::THttpRequestEnv& env) {
                AtomicIncrement(requestCount);
                return EchoReply(env);
            },
            {}
        );

        NJson::TJsonValue config;
        config["maxlen"] = 65536;
        config["maxreq"] = 65536;
        auto& dynamic = config["balancer"]["dynamic"];
        dynamic["max_pessimized_share"] = 0.5;

        NJson::TJsonValue proxy;
        proxy["cached_ip"] = "127.0.0.1";
        proxy["host"] = "localhost";
        proxy["port"] = echoService.Port;
        proxy["backend_timeout"] = "1s";

        NJson::TJsonValue item;
        item["proxy"] = proxy;
        dynamic.AppendValue(item);

        TConcurrentTest ct;
        ct.ScheduleRequests(env, config, false, true);

        while (AtomicGet(ct.ReqCounter) < 1000) {
            Sleep(TDuration::MilliSeconds(10));
        }

        AtomicSet(ct.Stopped, 1);
        ct.Stop();

        UNIT_ASSERT_VALUES_EQUAL(requestCount, ct.ReqCounter);    }

    Y_UNIT_TEST(ProxyClientWithAndWithoutSsl) {
        TEnv env;
        env.Start();

        TAtomic requestCount = 0;

        const TString sslConfigPath = "./ssl.cfg";
        {
            TFileOutput fo(sslConfigPath);
            MakeSslConfig(fo);
        }

        NBalancerServer::TOptions options;
        options.SslConfigPath = sslConfigPath;

        TService echoSslService([&requestCount](NBalancerServer::THttpRequestEnv &env) {
                AtomicIncrement(requestCount);
                return EchoReply(env);
            },
            options
        );

        auto port = echoSslService.Port;

        NJson::TJsonValue config;
        config["maxlen"] = 65536;
        config["maxreq"] = 65536;
        auto &rr = config["balancer"]["rr"];

        NJson::TJsonValue proxy;
        proxy["cached_ip"] = "127.0.0.1";
        proxy["host"] = "localhost";
        proxy["port"] = port;
        proxy["backend_timeout"] = "1s";
        proxy["https_settings"]["ciphers"] = "kEECDH:kRSA+AES128:kRSA:+3DES:RC4:!aNULL:!eNULL:!MD5:!EXPORT:!LOW:!SEED:!CAMELLIA:!IDEA:!PSK:!SRP:!SSLv2";
        proxy["https_settings"]["ca_file"] = CaPath;

        NJson::TJsonValue item;
        item["proxy"] = proxy;
        rr.AppendValue(item);

        TConcurrentTest ct;

        auto shoot = [&] (auto count, bool useEncryption) {
            AtomicSet(ct.Stopped, 0);

            ct.ScheduleRequests(env, config, useEncryption);

            while (AtomicGet(ct.ReqCounter) < count) {
                Sleep(TDuration::MilliSeconds(10));
            }

            AtomicSet(ct.Stopped, 1);
            ct.Stop();

            UNIT_ASSERT_VALUES_EQUAL(requestCount, ct.ReqCounter);
        };

        shoot(1000, /*useEncryption =*/ true);

        echoSslService.Server.Stop();

        TService echoService([&requestCount](NBalancerServer::THttpRequestEnv &env) {
                AtomicIncrement(requestCount);
                return EchoReply(env);
            },
            {},
            port
        );

        shoot(requestCount + 1000, /*useEncryption =*/ false);
    }

    Y_UNIT_TEST(SdClient) {
        TEnv env;
        env.Start();

        TAtomic requestCount1 = 0;
        TAtomic requestCount2 = 0;

        TService echoService1([&requestCount1](NBalancerServer::THttpRequestEnv& env) {
                AtomicIncrement(requestCount1);
                return EchoReply(env);
            },
            {}
        );

        TService echoService2([&requestCount2](NBalancerServer::THttpRequestEnv& env) {
                AtomicIncrement(requestCount2);
                return EchoReply(env);
            },
            {}
        );

        using namespace NYP::NServiceDiscovery;

        NApi::TReqResolveEndpoints req;
        req.set_cluster_name("test-cluster");
        req.set_endpoint_set_id("test-service");

        NJson::TJsonValue endpointSet;
        endpointSet["cluster_id"] = req.cluster_name();
        endpointSet["endpoint_set_id"] =  req.endpoint_set_id();

        NApi::TRspResolveEndpoints res1;
        {
            NApi::TEndpoint* endpoint = res1.mutable_endpoint_set()->add_endpoints();
            endpoint->set_fqdn("localhost");
            endpoint->set_port(echoService1.Port);
            endpoint->set_ip4_address("127.0.0.1");
            endpoint->set_ready(true);
        }

        NApi::TRspResolveEndpoints res2;
        {
            NApi::TEndpoint* endpoint = res2.mutable_endpoint_set()->add_endpoints();
            endpoint->set_fqdn("localhost");
            endpoint->set_port(echoService2.Port);
            endpoint->set_ip4_address("127.0.0.1");
            endpoint->set_ready(true);
        }

        env.SDServer().Set(req, res1);

        NJson::TJsonValue config;
        config["maxlen"] = 65536;
        config["maxreq"] = 65536;
        auto& sd = config["balancer"]["sd"];

        sd["proxy_options"]["backend_timeout"] = "10s";
        NJson::TJsonValue proxyWrapper;
        proxyWrapper["key"] = "dummy";
        proxyWrapper["value"] = NJson::JSON_MAP;
        sd["proxy_wrapper"] = proxyWrapper;
        sd["endpoint_sets"].AppendValue(endpointSet);
        sd["dynamic"] = NJson::TJsonMap();
        sd["dynamic"]["max_pessimized_share"] = 0.5;

        TConcurrentTest ct;
        ct.ScheduleRequests(env, config);

        for (size_t i = 0; i < 10; ++i) {
            while ((size_t)AtomicGet(ct.ReqCounter) < (i + 1) * 100) {
                Sleep(TDuration::MilliSeconds(10));
            }
            env.SDServer().Set(req, i % 2 ? res1 : res2);
        }

        AtomicSet(ct.Stopped, 1);
        ct.Stop();

        UNIT_ASSERT(requestCount1 > 0);
        UNIT_ASSERT(requestCount2 > 0);
        UNIT_ASSERT_VALUES_EQUAL(requestCount1 + requestCount2, ct.ReqCounter);
    }
}
