#include "server.h"

#include <balancer/server/ut/util/common.h>
#include <balancer/server/ut/util/env.h>
#include <balancer/server/ut/util/raw_socket_io.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/registar.h>

#include <library/cpp/cgiparam/cgiparam.h>
#include <library/cpp/string_utils/quote/quote.h>
#include <library/cpp/string_utils/base64/base64.h>

#include <util/system/fs.h>
#include <util/thread/pool.h>
#include <util/random/random.h>
#include <util/stream/file.h>

using namespace NBalancerServer;
using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TestServer) {
    Y_UNIT_TEST(Trivial) {
        TStandaloneServer server([](THttpRequestEnv&) {
                return TError();
            },
            TOptions().SetPort(80)
        );
    }

    Y_UNIT_TEST(Run) {
        NTesting::TEnv env;

        env.Start(NTesting::GetOkCallback(), TOptions());

        for (size_t i = 0; i < 100; ++i) {
            auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1")->Wait();
            UNIT_ASSERT(!res->IsError());
            UNIT_ASSERT_VALUES_EQUAL(res->Data, "OK");
        }

        env.Stop();
    }

    Y_UNIT_TEST(AccessLog) {
        NTesting::TEnv env;

        TOptions opts;
        opts.AccessLog = "./access_log";
        NFs::Remove(opts.AccessLog);

        env.Start(NTesting::GetOkCallback(), opts);

        for (size_t i = 0; i < 100; ++i) {
            auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1")->Wait();
            UNIT_ASSERT(!res->IsError());
            UNIT_ASSERT_VALUES_EQUAL(res->Data, "OK");
        }

        env.Stop();

        TFileInput fi(opts.AccessLog);
        TString accessLogLine;
        size_t lineCount = 0;
        while (fi.ReadLine(accessLogLine)) {
            ++lineCount;
        }

        UNIT_ASSERT_VALUES_EQUAL(lineCount, 100);
    }

    Y_UNIT_TEST(RequestData) {
        TInstant testStart = Now();
        auto checkRequest = [testStart](THttpRequestEnv& env) {
            UNIT_ASSERT_VALUES_EQUAL(env.FirstLine(), "GET /yandsearch?cgi=1 HTTP/1.1");
            UNIT_ASSERT(IsIn({"::1", "127.0.0.1"}, env.RemoteAddress()));
            UNIT_ASSERT(testStart < env.StartTime() && env.StartTime() < Now());


            NSrvKernel::TResponse response(200, "Ok");
            response.Props().ContentLength = 2;

            auto reply = env.GetReplyTransport();
            reply->SendHead(std::move(response));
            reply->SendData(TStringBuf("OK"));
            reply->SendEof();

            UNIT_ASSERT_VALUES_EQUAL(reply->SentSize(), 2);

            return TError();
        };

        NTesting::TEnv env;
        env.Start(checkRequest, TOptions());

        auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1")->Wait();
        UNIT_ASSERT(!res->IsError());
        UNIT_ASSERT_VALUES_EQUAL(res->Data, "OK");

        env.Stop();
    }

    Y_UNIT_TEST(Shutdown) {
        NTesting::TEnv testEnv;

        auto shutdownCallback = [&testEnv](THttpRequestEnv& env) {
            auto sendReply = [&](TStringBuf content) {
                NSrvKernel::TResponse response(200, "Ok");
                response.Props().ContentLength = content.Size() ;

                auto reply = env.GetReplyTransport();
                reply->SendHead(std::move(response));
                reply->SendData(ToString(content));
                reply->SendEof();
            };

            if (env.Path() == "/admin") {
                TQuickCgiParam cgi(env.Cgi());
                if (cgi.Has("action", "shutdown")) {
                    testEnv.Server->Shutdown();
                    sendReply("Stopped");
                    return TError();
                }
            }

            sendReply("OK");

            return TError();
        };

        testEnv.Start(shutdownCallback, TOptions());

        for (size_t i = 0; i < 100; ++i) {
            auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << testEnv.Port << "/yandsearch?cgi=1")->Wait();
            UNIT_ASSERT(!res->IsError());
            UNIT_ASSERT_VALUES_EQUAL(res->Data, "OK");
        }

        auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << testEnv.Port << "/admin?action=shutdown")->Wait();
        UNIT_ASSERT(!res->IsError());
        UNIT_ASSERT_VALUES_EQUAL(res->Data, "Stopped");

        for (size_t i = 1; i < 11; ++i) {
            if (!NCoro::TryConnect("localhost", testEnv.Port)) {
                return;
            }

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

        UNIT_ASSERT(!NCoro::TryConnect("localhost", testEnv.Port));
    }

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

    Y_UNIT_TEST(TestEcho) {
        for (size_t threads : {0, 1, 5}) {
            NTesting::RunConcurrent<NTesting::TEnv>(
                NTesting::GetEchoCallback(),
                [](NTesting::TEnv& env) {
                    const auto size = RandomNumber<size_t>(10 * 1024) + 1;
                    TString data = NUnitTest::RandomString(size, size);
                    auto msg = NNeh::TMessage::FromString(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1");

                    NNeh::NHttp::MakeFullRequest(msg, {}, data);
                    auto res = NNeh::Request(msg)->Wait();
                    UNIT_ASSERT(!res->IsError());
                    UNIT_ASSERT_VALUES_EQUAL(res->Data, data);
                },
                TOptions().SetThreads(threads)
            );
        }
    }

    auto echoCgiCallback = [](THttpRequestEnv& env) {
        NSrvKernel::TResponse response(200, "Ok");
        response.Props().ChunkedTransfer = true;

        TQuickCgiParam cgi(env.Cgi());

        auto reply = env.GetReplyTransport();
        reply->SendHead(std::move(response));
        reply->SendData(CGIEscapeRet(cgi.Get("reply")));
        reply->SendEof();

        return TError();
    };

    Y_UNIT_TEST(TestEchoCgi) {
        for (size_t threads : {0, 1, 5}) {
            NTesting::RunConcurrent<NTesting::TEnv>(
                echoCgiCallback,
                [](NTesting::TEnv& env) {
                    const auto size = RandomNumber<size_t>(10 * 1024) + 1;
                    TString data = NUnitTest::RandomString(size, size);
                    CGIEscape(data);
                    auto msg = NNeh::TMessage::FromString(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?reply=" << data);

                    auto res = NNeh::Request(msg)->Wait();
                    UNIT_ASSERT(!res->IsError());
                    UNIT_ASSERT_VALUES_EQUAL(res->Data, data);
                },
                TOptions().SetThreads(threads)
            );
        }
    }

    auto echoHeadersCallback = [](THttpRequestEnv& env) {
        NSrvKernel::TResponse response(200, "Ok");
        response.Props().ChunkedTransfer = true;

        TString data = TString(env.Headers().GetFirstValue("Reply"));
        response.Headers().Add("Reply", data);

        auto reply = env.GetReplyTransport();
        reply->SendHead(std::move(response));
        reply->SendData(data);
        reply->SendEof();

        return TError();
    };

    Y_UNIT_TEST(TestEchoHeaders) {
        for (size_t threads : {0, 1, 5}) {
            NTesting::RunConcurrent<NTesting::TEnv>(
                echoHeadersCallback,
                [](NTesting::TEnv& env) {
                    const auto size = RandomNumber<size_t>(10 * 1024) + 1;
                    TString data = Base64Encode(NUnitTest::RandomString(size, size));

                    auto msg = NNeh::TMessage::FromString(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1");


                    NNeh::NHttp::MakeFullRequest(msg, TStringBuilder() << "Reply: " << data << "\r\n", "");
                    auto res = NNeh::Request(msg)->Wait();
                    UNIT_ASSERT(!res->IsError());
                    UNIT_ASSERT_VALUES_EQUAL(res->Data, data);

                    auto* header = res->Headers.FindHeader("Reply");
                    UNIT_ASSERT(header);
                    UNIT_ASSERT_VALUES_EQUAL(header->Value(), data);
                },
                TOptions().SetThreads(threads)
            );
        }
    }

    void RunEnqueueTest(size_t threads, size_t queueSize) {
        NTesting::TEnv env;

        TOptions options;
        options.SetNetworkThreads(5).SetThreads(threads).SetQueueSize(queueSize);
        const size_t serverCapacity = threads + queueSize;

        TManualEvent event;

        auto delayedCallback = [&event](THttpRequestEnv& request) {
            event.WaitI();
            return NTesting::GetOkCallback()(request);
        };

        env.Start(delayedCallback, options);

        for (size_t run = 0; run < 3; ++run) {
            TAtomic succCounter = 0;
            TAtomic failCounter = 0;

            TThreadPool shoot;
            shoot.Start(serverCapacity + 5, 0);
            const size_t extraRequests = 5 + 5 * run;
            for (size_t i = 0; i < serverCapacity + extraRequests; ++i) {
                UNIT_ASSERT(shoot.AddFunc(
                    [&, i]() {
                        auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=" << i)->Wait();
                        if (res->IsError()) {
                            UNIT_ASSERT_VALUES_EQUAL(res->GetErrorCode(), 503);
                            AtomicIncrement(failCounter);
                        } else {
                            UNIT_ASSERT_VALUES_EQUAL(res->Data, "OK");
                            AtomicIncrement(succCounter);
                        }
                    }
                ));
            }

            const TInstant start = Now();
            while ((Now() - start) < TDuration::Seconds(10) && (size_t)AtomicGet(failCounter) < extraRequests) {
                Sleep(TDuration::MilliSeconds(100));
            }

            event.Signal();
            shoot.Stop();
            UNIT_ASSERT_VALUES_EQUAL(succCounter, serverCapacity);
            UNIT_ASSERT_VALUES_EQUAL(failCounter, extraRequests);
            event.Reset();
        }

        env.Stop();
    }


    Y_UNIT_TEST(TestEnqueue_1_1) {
        RunEnqueueTest(1, 1);
    }

    Y_UNIT_TEST(TestEnqueue_1_10) {
        RunEnqueueTest(1, 10);
    }

    Y_UNIT_TEST(TestEnqueue_10_10) {
        RunEnqueueTest(10, 10);
    }

    Y_UNIT_TEST(TestExceptions) {
        for (size_t threads : {0, 1, 5}) {
            TAtomic succ = 0, fail = 0;

            NTesting::RunConcurrent<NTesting::TEnv>(
                [](THttpRequestEnv& env) {
                    NSrvKernel::TResponse response(200, "Ok");
                    response.Props().ChunkedTransfer = true;

                    TQuickCgiParam cgi(env.Cgi());
                    const bool makeException = cgi.Has("exception", "yes");

                    TString body = env.Body();

                    if (makeException && RandomNumber<double>() < 0.5) {
                        ythrow yexception() << "error";
                    }

                    auto reply = env.GetReplyTransport();
                    reply->SendHead(std::move(response));

                    size_t pos = 0;
                    while (pos < body.Size()) {
                        if (makeException && RandomNumber<double>() < 0.1) {
                            ythrow yexception() << "error";
                        }

                        size_t next = Min<size_t>(pos + 10, body.Size());
                        TString chunk = body.substr(pos, next - pos);
                        reply->SendData(chunk);
                        pos = next;
                    }

                    if (makeException) {
                        ythrow yexception() << "error";
                    }

                    reply->SendEof();

                    return TError();
                },
                [&succ, &fail](NTesting::TEnv& env) {
                    TStringBuilder request;
                    request << "http://localhost:" << env.Port << "/yandsearch?cgi=1";
                    const bool makeException = RandomNumber<double>() < 0.9;
                    if (makeException) {
                        request << "&exception=yes";
                    }

                    auto msg = NNeh::TMessage::FromString(request);
                    const auto size = RandomNumber<size_t>(10 * 1024) + 1;
                    TString data = Base64Encode(NUnitTest::RandomString(size, size));
                    NNeh::NHttp::MakeFullRequest(msg, {}, data);

                    auto res = NNeh::Request(msg)->Wait();
                    if (makeException) {
                        AtomicIncrement(fail);
                        UNIT_ASSERT(res->IsError());
                        UNIT_ASSERT_VALUES_EQUAL(res->GetErrorText(), "Connection reset by peer");
                    } else {
                        AtomicIncrement(succ);
                        UNIT_ASSERT(!res->IsError());
                        UNIT_ASSERT_VALUES_EQUAL(res->Data, data);
                    }
                },
                TOptions().SetThreads(threads)
            );

            UNIT_ASSERT(AtomicGet(succ) > 0);
            UNIT_ASSERT(AtomicGet(fail) > 0);
        }
    }

    Y_UNIT_TEST(TestSecondaryThreadPool) {
        TThreadPool pool;
        pool.Start(10);

        for (size_t threads : {0, 1, 5}) {
            NTesting::RunConcurrent<NTesting::TEnv>(
                [&pool](THttpRequestEnv& env) {
                    UNIT_ASSERT(
                        pool.AddFunc([reply = env.GetReplyTransport(), body = env.Body()]() {
                            NSrvKernel::TResponse response(200, "Ok");
                            response.Props().ChunkedTransfer = true;

                            reply->SendHead(std::move(response));
                            reply->SendData(body);
                            reply->SendEof();
                        })
                    );

                    return TError();
                },
                [](NTesting::TEnv& env) {
                    const auto size = RandomNumber<size_t>(10 * 1024) + 1;
                    TString data = NUnitTest::RandomString(size, size);
                    auto msg = NNeh::TMessage::FromString(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1");

                    NNeh::NHttp::MakeFullRequest(msg, {}, data);
                    auto res = NNeh::Request(msg)->Wait();
                    UNIT_ASSERT(!res->IsError());
                    UNIT_ASSERT_VALUES_EQUAL(res->Data, data);
                },
                TOptions().SetThreads(threads)
            );
        }

        pool.Stop();
    }

    Y_UNIT_TEST(TestConnectionHeader) {
        auto cb = [](THttpRequestEnv& env) {
            UNIT_ASSERT(env.Headers().GetFirstValue("Connection") == "Keep-Alive");

            NSrvKernel::TResponse response(200, "Ok");
            response.Props().ChunkedTransfer = true;
            response.Props().AddExplicitConnectionHeader = true;

            auto reply = env.GetReplyTransport();
            reply->SendHead(std::move(response));
            reply->SendEof();

            return TError();
        };

        NTesting::TEnv env;

        env.Start(cb, TOptions());

        auto f = [&] (TCont* cont) {
            try {
                NTesting::TRawSocketIO io("localhost", env.Port, cont, Now() + TDuration::Seconds(10));

                THttpOutput out(&io);
                out.EnableKeepAlive(true);
                out << "GET /?xxx=1 HTTP/1.1\r\n\r\n";
                out.Flush();

                io.Flush();

                THttpInput httpI(&io);

                UNIT_ASSERT_VALUES_EQUAL(ParseHttpRetCode(httpI.FirstLine()), 200);

                auto* connectionHeader = httpI.Headers().FindHeader("Connection");
                UNIT_ASSERT(connectionHeader && connectionHeader->Value() == "Keep-Alive");
                httpI.ReadAll();
            } catch (...) {
                UNIT_ASSERT_C(false, CurrentExceptionMessage());
            }
        };

        TContExecutor e(128 * 1024, IPollerFace::Default(), nullptr, nullptr, NCoro::NStack::EGuard::Canary,
                        NCoro::NStack::TPoolAllocatorSettings{});
        e.Create(f, "request");
        e.Execute();

        env.Stop();
    }

    Y_UNIT_TEST(TestClientError) {
        auto cb = [](THttpRequestEnv& env) {

            NSrvKernel::TResponse response(200, "Ok");
            response.Props().ChunkedTransfer = true;
            response.Props().AddExplicitConnectionHeader = true;


            TString failed;

            auto reply = env.GetReplyTransport();
            try {
                reply->SendHead(std::move(response));

                for (size_t i = 0; i < 10000; ++i) {
                    reply->SendData(TString("xxx"));
                }

                reply->SendEof();
            } catch (...) {
                failed = CurrentExceptionMessage();
            }

            UNIT_ASSERT_STRING_CONTAINS(failed, "IHttpReplyTransport::TClientError");

            return TError();
        };

        NTesting::TEnv env;

        env.Start(cb, TOptions());

        auto f = [&] (TCont* cont) {
            try {
                NTesting::TRawSocketIO io("localhost", env.Port, cont, Now() + TDuration::Seconds(10));

                THttpOutput out(&io);
                out.EnableKeepAlive(true);
                out << "GET /?xxx=1 HTTP/1.1\r\n\r\n";
                out.Flush();

                io.Flush();
            } catch (...) {
                UNIT_ASSERT_C(false, CurrentExceptionMessage());
            }
        };

        TContExecutor e(128 * 1024, IPollerFace::Default(), nullptr, nullptr, NCoro::NStack::EGuard::Canary,
                        NCoro::NStack::TPoolAllocatorSettings{});
        e.Create(f, "request");
        e.Execute();

        env.Stop();
    }

    Y_UNIT_TEST(SilentCallback1) {
        auto silentCb = [](THttpRequestEnv&) {
            return TError();
        };

        for (size_t threads : {0, 1, 5}) {
            NTesting::TEnv env;
            env.Start(silentCb, TOptions().SetThreads(threads));

            auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1")->Wait(TDuration::Seconds(5));
            UNIT_ASSERT(res);
            UNIT_ASSERT(res->IsError());

            env.Stop();
        }
    }

    Y_UNIT_TEST(SilentCallback2) {
        auto silentCb = [](THttpRequestEnv& env) {
            env.GetReplyTransport();
            return TError();
        };

        for (size_t threads : {0, 1, 5}) {
            NTesting::TEnv env;
            env.Start(silentCb, TOptions().SetThreads(threads));

            auto res = NNeh::Request(TStringBuilder() << "http://localhost:" << env.Port << "/yandsearch?cgi=1")->Wait(TDuration::Seconds(5));
            UNIT_ASSERT(res);
            UNIT_ASSERT(res->IsError());
            env.Stop();
        }

    }
}

