/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.celeborn.common.network.protocol;

import java.io.EOFException;
import java.io.InputStream;

import javax.annotation.Nullable;

import com.google.common.base.Preconditions;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.stream.ChunkedInput;
import io.netty.handler.stream.ChunkedStream;
import io.netty.util.ReferenceCountUtil;

import org.apache.celeborn.common.network.buffer.ManagedBuffer;

/**
 * A wrapper message that holds two separate pieces (a header and a body).
 *
 * <p>The header must be a ByteBuf, while the body can be a ByteBuf, InputStream, or ChunkedStream.
 * Based on common/network-common/org.apache.spark.network.protocol.EncryptedMessageWithHeader
 */
public class EncryptedMessageWithHeader implements ChunkedInput<ByteBuf> {

  @Nullable private final ManagedBuffer managedBuffer;
  private final ByteBuf header;
  private final int headerLength;
  private final Object body;
  private final long bodyLength;
  private long totalBytesTransferred;

  /**
   * Construct a new EncryptedMessageWithHeader.
   *
   * @param managedBuffer the {@link ManagedBuffer} that the message body came from. This needs to
   *     be passed in so that the buffer can be freed when this message is deallocated. Ownership of
   *     the caller's reference to this buffer is transferred to this class, so if the caller wants
   *     to continue to use the ManagedBuffer in other messages then they will need to call retain()
   *     on it before passing it to this constructor.
   * @param header the message header.
   * @param body the message body.
   * @param bodyLength the length of the message body, in bytes.
   */
  public EncryptedMessageWithHeader(
      @Nullable ManagedBuffer managedBuffer, ByteBuf header, Object body, long bodyLength) {
    Preconditions.checkArgument(
        body instanceof ByteBuf || body instanceof InputStream || body instanceof ChunkedStream,
        "Body must be a ByteBuf, an InputStream, or a ChunkedStream.");
    this.managedBuffer = managedBuffer;
    this.header = header;
    this.headerLength = header.readableBytes();
    this.body = body;
    this.bodyLength = bodyLength;
    this.totalBytesTransferred = 0;
  }

  @Override
  public ByteBuf readChunk(ChannelHandlerContext ctx) throws Exception {
    return readChunk(ctx.alloc());
  }

  @Override
  public ByteBuf readChunk(ByteBufAllocator allocator) throws Exception {
    if (isEndOfInput()) {
      return null;
    }

    if (body instanceof ByteBuf) {
      // For ByteBuf bodies, return header + body as a single composite buffer.
      ByteBuf bodyBuf = (ByteBuf) body;
      totalBytesTransferred = headerLength + bodyLength;
      return Unpooled.wrappedBuffer(header.retain(), bodyBuf.retain());
    } else {
      if (totalBytesTransferred < headerLength) {
        totalBytesTransferred += headerLength;
        return header.retain();
      } else if (body instanceof InputStream) {
        InputStream stream = (InputStream) body;
        int available = stream.available();
        if (available <= 0) {
          available = (int) (length() - totalBytesTransferred);
        } else {
          available = (int) Math.min(available, length() - totalBytesTransferred);
        }
        ByteBuf buffer = allocator.buffer(available);
        int toRead = Math.min(available, buffer.writableBytes());
        int read = buffer.writeBytes(stream, toRead);
        if (read >= 0) {
          totalBytesTransferred += read;
          return buffer;
        } else {
          throw new EOFException("Unable to read bytes from InputStream");
        }
      } else if (body instanceof ChunkedStream) {
        ChunkedStream stream = (ChunkedStream) body;
        long old = stream.transferredBytes();
        ByteBuf buffer = stream.readChunk(allocator);
        long read = stream.transferredBytes() - old;
        if (read >= 0) {
          totalBytesTransferred += read;
          assert (totalBytesTransferred <= length());
          return buffer;
        } else {
          throw new EOFException("Unable to read bytes from ChunkedStream");
        }
      } else {
        return null;
      }
    }
  }

  @Override
  public long length() {
    return headerLength + bodyLength;
  }

  @Override
  public long progress() {
    return totalBytesTransferred;
  }

  @Override
  public boolean isEndOfInput() throws Exception {
    return (headerLength + bodyLength) == totalBytesTransferred;
  }

  @Override
  public void close() throws Exception {
    header.release();
    ReferenceCountUtil.release(body);
    if (managedBuffer != null) {
      managedBuffer.release();
    }
    if (body instanceof InputStream) {
      ((InputStream) body).close();
    } else if (body instanceof ChunkedStream) {
      ((ChunkedStream) body).close();
    }
  }
}
