#include "transport.h"

#include <library/cpp/http/io/compression.h>
#include <library/cpp/http/io/headers.h>

#include <util/string/ascii.h>
#include <util/string/strip.h>

namespace NBalancerServer {
    IHttpReplyTransport::~IHttpReplyTransport() {}

    class IHttpReplyTransport::TSendDataStream
        : public IOutputStream
    {
    public:
        explicit TSendDataStream(IHttpReplyTransport& transport)
            : Transport_(transport)
        {
        }

        void DoWrite(const void* data, size_t size) override {
            TString str((const char*)data, size);
            Transport_.SendData(NSrvKernel::TChunkList(std::move(str)));
        }

    private:
        IHttpReplyTransport& Transport_;
    };

    bool IHttpReplyTransport::ChooseEncoding(const NSrvKernel::THeaders& requestHeaders) {
        Y_ENSURE(!SendState_.HeadSent, "try to set encoding after send head");
        Y_ENSURE(!Encoder_, "try to setup encoder twice");

        THashSet<TStringBuf, TCIOps, TCIOps> accepted;
        for (const auto& value : requestHeaders.GetValuesRef("Accept-Encoding")) {
            for (const auto& it : StringSplitter(value.AsStringBuf()).Split(',').SkipEmpty()) {
                accepted.insert(StripString(it.Token()));
            }
        }

        Encoding_ = NHttp::ChooseBestCompressionScheme(
            [&accepted](TStringBuf v) {
                return accepted.contains(v);
            },
            TCompressionCodecFactory::Instance().GetBestCodecs()
        );

        return !Encoding_.Empty();
    }

    void IHttpReplyTransport::SendHead(NSrvKernel::TResponse&& response) {
        CheckClientError();
        CheckNotFinished();

        Y_ENSURE(!SendState_.HeadSent, "head already sent");
        SendState_.HeadSent = true;

        if (Encoding_ && response.Props().ChunkedTransfer && !response.Headers().FindValues("Content-Encoding")) {
            if (const auto* encoderFactory = TCompressionCodecFactory::Instance().FindEncoder(Encoding_)) {
                SendDataStream_ = MakeHolder<TSendDataStream>(*this);
                Encoder_ = (*encoderFactory)(SendDataStream_.Get());
                response.Headers().Add("Content-Encoding", Encoding_);
            }
        }

        DoSendHead(std::move(response));
    }

    // Call and check NBalancerServer::THttpRequestEnv::TrailersAllowed() before using that method
    void IHttpReplyTransport::SendTrailers(NSrvKernel::THeaders&& headers) {
        CheckClientError();
        CheckNotFinished();

        DoSendTrailers(std::move(headers));
    }
}
