#include <balancer/kernel/custom_io/chunkedtransferio.h>
#include <balancer/kernel/custom_io/chunkio.h>
#include <balancer/kernel/http/parser/httpencoder.h>
#include <balancer/kernel/http/parser/output.h>
#include <balancer/kernel/http/parser/response_builder.h>

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

using namespace NSrvKernel;

Y_UNIT_TEST_SUITE(TToBackendEncoderTest) {
    Y_UNIT_TEST(TestDuplicateHeadersBug) {
        const TString bodyStr("xxxxx");
        TRequest req;
        {
            TChunkList headers("GET / HTTP/1.1\r\nHost: yandex.ru\r\nContent-Length: 5\r\n\r\n");
            TChunksInput headersIn(std::move(headers));
            TChunkList lst;
            auto err = req.Read(&headersIn, lst, TInstant::Max());
            UNIT_ASSERT(!err);
        }
        {
            TChunksOutput out;
            TToBackendEncoder enc(&out);

            auto err = enc.WriteRequest(req, false, Nothing(), TInstant::Max());
            UNIT_ASSERT(!err);

            auto rawResponse = ToString(out.Chunks());

            UNIT_ASSERT(rawResponse.StartsWith("GET / HTTP/1.1\r\n"));
            UNIT_ASSERT_UNEQUAL(rawResponse.find("Host: yandex.ru\r\n"), TString::npos);
            UNIT_ASSERT_UNEQUAL(rawResponse.find("Content-Length: 5\r\n"), TString::npos);
            UNIT_ASSERT_UNEQUAL(rawResponse.find("Connection: Close\r\n"), TString::npos);

            out.Chunks().Clear();
            err = enc.Send(TChunkList(bodyStr), TInstant::Max());
            UNIT_ASSERT(!err);
            err = enc.SendEof(TInstant::Max());
            UNIT_ASSERT(!err);

            rawResponse = ToString(out.Chunks());
            UNIT_ASSERT_EQUAL(rawResponse, "xxxxx");
        }
        // BALANCER-2146 Yes, TToBackendEncoder used to add extra headers to TRequest passed to it. /O
        // Another great piece of code sitting in the plain sight for years.
        {
            TChunksOutput out;
            TToBackendEncoder enc(&out);

            auto err = enc.WriteRequest(req, true, Nothing(), TInstant::Max());
            UNIT_ASSERT(!err);

            auto rawResponse = ToString(out.Chunks());

            UNIT_ASSERT(rawResponse.StartsWith("GET / HTTP/1.1\r\n"));
            UNIT_ASSERT_UNEQUAL(rawResponse.find("Host: yandex.ru\r\n"), TString::npos);
            UNIT_ASSERT_UNEQUAL(rawResponse.find("Content-Length: 5\r\n"), TString::npos);
            UNIT_ASSERT_EQUAL(rawResponse.find("Connection: Close\r\n"), TString::npos);

            out.Chunks().Clear();
            err = enc.Send(TChunkList(bodyStr), TInstant::Max());
            UNIT_ASSERT(!err);
            err = enc.SendEof(TInstant::Max());
            UNIT_ASSERT(!err);

            rawResponse = ToString(out.Chunks());
            UNIT_ASSERT_EQUAL(rawResponse, "xxxxx");
        }
    }
}

Y_UNIT_TEST_SUITE(TToClientEncoderTest) {
    const TString BODY_SAMPLE = TString("Lorem ipsum dolor sit amet, consectetur adipiscing"
                                             " elit, sed do eiusmod tempor incididunt ut labore et"
                                             " dolore magna aliqua. Ut enim ad minim veniam, quis"
                                             " nostrud exercitation ullamco laboris nisi ut aliquip"
                                             " ex ea commodo consequat. Duis aute irure dolor in"
                                             " reprehenderit in voluptate velit esse cillum dolore"
                                             " eu fugiat nulla pariatur. Excepteur sint occaecat"
                                             " cupidatat non proident, sunt in culpa qui officia"
                                             " deserunt mollit anim id est laborum.");

    Y_UNIT_TEST(Test11To11Chunked200) {
        TChunksIo chunksIo;
        THttpOutput httpOut{&chunksIo, 1, false, false, true};
        TResponse response = BuildResponse().Version11()
                                            .Code(HTTP_OK)
                                            .Reason("OK")
                                            .ChunkedTransfer();

        TryRethrowError(httpOut.SendHead(std::move(response), false, TInstant::Max()));
        auto rawResponse = ToString(chunksIo.Chunks());
        UNIT_ASSERT(rawResponse.StartsWith("HTTP/1.1 200 OK\r\n"));
        UNIT_ASSERT_UNEQUAL(rawResponse.find("Transfer-Encoding: chunked\r\n"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Connection"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Content-Length"), TString::npos);

        chunksIo.Chunks().Clear();
        TryRethrowError(httpOut.Send(TChunkList{BODY_SAMPLE}, TInstant::Max()));
        TryRethrowError(httpOut.SendEof(TInstant::Max()));

        TChunkedTransferDecoder chunksIn{&chunksIo};
        TChunkList rawBody;
        TryRethrowError(RecvAll(&chunksIn, rawBody, TInstant::Max()));
        UNIT_ASSERT_EQUAL(rawBody, BODY_SAMPLE);
    }

    Y_UNIT_TEST(Test11To11ContentLength200) {
        TChunksIo chunksIo;
        THttpOutput httpOut{&chunksIo, 1, false, false, true};
        TResponse response = BuildResponse().Version11()
                                            .Code(HTTP_OK)
                                            .Reason("OK")
                                            .ContentLength(BODY_SAMPLE.size());

        TryRethrowError(httpOut.SendHead(std::move(response), false, TInstant::Max()));
        auto rawResponse = ToString(chunksIo.Chunks());
        UNIT_ASSERT(rawResponse.StartsWith("HTTP/1.1 200 OK\r\n"));
        UNIT_ASSERT_UNEQUAL(rawResponse.find("Content-Length: " + ToString(BODY_SAMPLE.size())
                                                                + "\r\n"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Connection"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Transfer-Encoding"), TString::npos);

        chunksIo.Chunks().Clear();
        TryRethrowError(httpOut.Send(TChunkList{BODY_SAMPLE}, TInstant::Max()));
        TryRethrowError(httpOut.SendEof(TInstant::Max()));

        UNIT_ASSERT_EQUAL(chunksIo.Chunks(), BODY_SAMPLE);
    }

    Y_UNIT_TEST(TestUpgradeRequested) {
        TChunksIo chunksIo;
        THttpOutput httpOut{&chunksIo, 1, false, false, true};
        TResponse response = BuildResponse().Version11()
                                            .Code(HTTP_SWITCHING_PROTOCOLS)
                                            .Reason("Switching protocols")
                                            .UpgradeRequested();

        TryRethrowError(httpOut.SendHead(std::move(response), false, TInstant::Max()));
        auto rawResponse = ToString(chunksIo.Chunks());
        UNIT_ASSERT(rawResponse.StartsWith("HTTP/1.1 101 Switching protocols\r\n"));
        UNIT_ASSERT_UNEQUAL(rawResponse.find("Connection: Upgrade\r\n"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Content-Length"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Transfer-Encoding"), TString::npos);

        chunksIo.Chunks().Clear();
        TryRethrowError(httpOut.Send(TChunkList{BODY_SAMPLE}, TInstant::Max()));
        TryRethrowError(httpOut.SendEof(TInstant::Max()));

        UNIT_ASSERT_EQUAL(chunksIo.Chunks(), BODY_SAMPLE);
    }

    Y_UNIT_TEST(TestUpgradeRequestedWithContentLength) {
        TChunksIo chunksIo;
        THttpOutput httpOut{&chunksIo, 1, false, false, true};
        TResponse response = BuildResponse().Version11()
                                            .Code(HTTP_SWITCHING_PROTOCOLS)
                                            .Reason("Switching protocols")
                                            .UpgradeRequested();
        response.Props().ContentLength = 0;

        TryRethrowError(httpOut.SendHead(std::move(response), false, TInstant::Max()));
        auto rawResponse = ToString(chunksIo.Chunks());
        UNIT_ASSERT(rawResponse.StartsWith("HTTP/1.1 101 Switching protocols\r\n"));
        UNIT_ASSERT_UNEQUAL(rawResponse.find("Connection: Upgrade\r\n"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Content-Length"), TString::npos);
        UNIT_ASSERT_EQUAL(rawResponse.find("Transfer-Encoding"), TString::npos);

        chunksIo.Chunks().Clear();
        TryRethrowError(httpOut.Send(TChunkList{BODY_SAMPLE}, TInstant::Max()));
        TryRethrowError(httpOut.SendEof(TInstant::Max()));

        UNIT_ASSERT_EQUAL(chunksIo.Chunks(), BODY_SAMPLE);
    }
}
