/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package com.facebook.thrift.rsocket.client;

import com.facebook.thrift.client.RpcClient;
import com.facebook.thrift.client.RpcOptions;
import com.facebook.thrift.model.StreamResponse;
import com.facebook.thrift.payload.ClientRequestPayload;
import com.facebook.thrift.payload.ClientResponsePayload;
import com.facebook.thrift.payload.Reader;
import com.facebook.thrift.protocol.ByteBufTProtocol;
import com.facebook.thrift.protocol.TProtocolType;
import com.facebook.thrift.util.resources.RpcResources;
import io.netty.buffer.ByteBuf;
import io.netty.util.ReferenceCountUtil;
import java.util.function.Function;
import org.apache.thrift.ClientPushMetadata;
import org.reactivestreams.Publisher;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

public final class HeaderAwareRSocketRpcClient implements RpcClient {

  private static final String HEADER_KEY = "header_response";

  private final RpcClient delegate;

  public HeaderAwareRSocketRpcClient(RpcClient delegate) {
    this.delegate = delegate;
  }

  @Override
  public Mono<Void> onClose() {
    return delegate.onClose();
  }

  @Override
  public void close() {
    delegate.close();
  }

  @Override
  public <T> Mono<ClientResponsePayload<T>> singleRequestSingleResponse(
      ClientRequestPayload<T> payload, RpcOptions options) {
    return delegate.singleRequestSingleResponse(payload, options);
  }

  @Override
  public <T> Mono<Void> singleRequestNoResponse(
      ClientRequestPayload<T> payload, RpcOptions options) {
    return delegate.singleRequestNoResponse(payload, options);
  }

  @Override
  public Mono<Void> metadataPush(ClientPushMetadata clientMetadata, RpcOptions options) {
    return delegate.metadataPush(clientMetadata, options);
  }

  @Override
  public <T, K> Flux<ClientResponsePayload<K>> singleRequestStreamingResponse(
      ClientRequestPayload<T> payload, RpcOptions options) {
    return delegate
        .<T, K>singleRequestStreamingResponse(payload, options)
        .map(new HeaderResponseHandler<>(payload.getResponseReader()));
  }

  @Override
  public <T, K> Flux<ClientResponsePayload<K>> streamingRequestStreamingResponse(
      Publisher<ClientRequestPayload<T>> payloads, RpcOptions options) {
    return delegate.streamingRequestStreamingResponse(payloads, options);
  }

  /**
   * Converts header ClientResponsePayload<StreamResponse<Object>> to
   * ClientResponsePayload<StreamResponse<k>>
   *
   * @param <T> Generic type of request object
   * @param <K> Generic type of response object
   */
  private static class HeaderResponseHandler<T, K>
      implements Function<ClientResponsePayload<K>, ClientResponsePayload<K>> {

    private final Reader<T> responseReader;

    private HeaderResponseHandler(Reader<T> responseReader) {
      this.responseReader = responseReader;
    }

    @Override
    @SuppressWarnings({"rawtypes", "unchecked"})
    public ClientResponsePayload<K> apply(ClientResponsePayload<K> kClientResponsePayload) {
      StreamResponse response = ((StreamResponse) kClientResponsePayload.getData());
      if (response != null
          && response.isSetData()
          && response.getData() instanceof String
          && response.getData().equals(HEADER_KEY)) {

        ByteBuf buffer = RpcResources.getByteBufAllocator().buffer();
        try {
          ByteBufTProtocol dataProtocol = TProtocolType.TBinary.apply(buffer);
          K streamResponse = (K) StreamResponse.fromData(responseReader.read(dataProtocol));

          return ClientResponsePayload.createStreamResult(
              streamResponse,
              kClientResponsePayload.getResponseRpcMetadata(),
              kClientResponsePayload.getStreamPayloadMetadata(),
              kClientResponsePayload.getBinaryHeaders(),
              kClientResponsePayload.getStreamId());
        } finally {
          ReferenceCountUtil.safeRelease(buffer);
        }
      }
      return kClientResponsePayload;
    }
  }
}

@@ -996,204 +996,54 @@
 col.
-ByteBufTProtocol;%0Aimport com.facebook.thrift.protocol.TProtocolType;%0Aimport com.facebook.thrift.util.resources.RpcResources;%0Aimport io.netty.buffer.ByteBuf;%0Aimport io.netty.util.ReferenceCountUtil
+TProtocolType;%0Aimport io.netty.buffer.Unpooled
 ;%0Aim
@@ -3710,195 +3710,58 @@
     
-ByteBuf buffer = RpcResources.getByteBufAllocator().buffer();%0A        try %7B%0A          ByteBufTProtocol dataProtocol = TProtocolType.TBinary.apply(buffer);%0A          K streamResponse = (K)
+K streamResponse =%0A            (K)%0A               
  Str
@@ -3781,16 +3781,37 @@
 romData(
+%0A                    
 response
@@ -3826,20 +3826,17 @@
 ead(
-data
+T
 Protocol
 ));%0A
@@ -3835,15 +3835,54 @@
 ocol
+Type.TBinary.apply(Unpooled.EMPTY_BUFFER)
 ));%0A%0A
-  
     
@@ -3946,18 +3946,16 @@
         
-  
 streamRe
@@ -3962,18 +3962,16 @@
 sponse,%0A
-  
         
@@ -4027,34 +4027,32 @@
 (),%0A            
- 
- 
 kClientResponseP
@@ -4090,34 +4090,32 @@
 (),%0A            
-  
 kClientResponseP
@@ -4137,26 +4137,24 @@
 yHeaders(),%0A
-  
             
@@ -4196,88 +4196,8 @@
 ));%0A
-        %7D finally %7B%0A          ReferenceCountUtil.safeRelease(buffer);%0A        %7D%0A
     

use unpooled byebuf empty