/*
 * 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.nio.ByteBuffer;
import java.util.Objects;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;

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

/** An on-the-wire transmittable message. */
public abstract class Message implements Encodable {
  private ManagedBuffer body;

  protected Message() {
    this(null);
  }

  protected Message(ManagedBuffer body) {
    this.body = body;
  }

  /** Used to identify this request type. */
  public abstract Type type();

  /** An optional body for the message. */
  public ManagedBuffer body() {
    return body;
  }

  public void setBody(ByteBuf buf) {
    this.body = new NettyManagedBuffer(buf);
  }

  /** Whether the body should be copied out in frame decoder. */
  public boolean needCopyOut() {
    return false;
  }

  @SuppressWarnings("NonOverridingEquals")
  protected boolean equals(Message other) {
    return Objects.equals(body, other.body);
  }

  public ByteBuffer toByteBuffer() {
    // Allow room for encoded message, plus the type byte
    ByteBuf buf = Unpooled.buffer(encodedLength() + 1);
    buf.writeByte(type().id());
    encode(buf);
    assert buf.writableBytes() == 0 : "Writable bytes remain: " + buf.writableBytes();
    return buf.nioBuffer();
  }

  /** Preceding every serialized Message is its type, which allows us to deserialize it. */
  public enum Type implements Encodable {
    UNKNOWN_TYPE(-1),
    CHUNK_FETCH_REQUEST(0),
    CHUNK_FETCH_SUCCESS(1),
    CHUNK_FETCH_FAILURE(2),
    RPC_REQUEST(3),
    RPC_RESPONSE(4),
    RPC_FAILURE(5),
    OPEN_STREAM(6),
    STREAM_HANDLE(7),
    ONE_WAY_MESSAGE(9),
    PUSH_DATA(11),
    PUSH_MERGED_DATA(12),
    REGION_START(13),
    REGION_FINISH(14),
    PUSH_DATA_HAND_SHAKE(15),
    READ_ADD_CREDIT(16),
    READ_DATA(17),
    OPEN_STREAM_WITH_CREDIT(18),
    BACKLOG_ANNOUNCEMENT(19),
    TRANSPORTABLE_ERROR(20),
    BUFFER_STREAM_END(21),
    HEARTBEAT(22),
    SEGMENT_START(23),
    NOTIFY_REQUIRED_SEGMENT(24),
    SUBPARTITION_READ_DATA(25);

    private final byte id;

    Type(int id) {
      assert id < 128 : "Cannot have more than 128 message types";
      this.id = (byte) id;
    }

    public byte id() {
      return id;
    }

    @Override
    public int encodedLength() {
      return 1;
    }

    @Override
    public void encode(ByteBuf buf) {
      buf.writeByte(id);
    }

    public static Type decode(ByteBuffer buffer) {
      ByteBuf buf = Unpooled.wrappedBuffer(buffer);
      return decode(buf);
    }

    public static Type decode(ByteBuf buf) {
      byte id = buf.readByte();
      switch (id) {
        case 0:
          return CHUNK_FETCH_REQUEST;
        case 1:
          return CHUNK_FETCH_SUCCESS;
        case 2:
          return CHUNK_FETCH_FAILURE;
        case 3:
          return RPC_REQUEST;
        case 4:
          return RPC_RESPONSE;
        case 5:
          return RPC_FAILURE;
        case 6:
          return OPEN_STREAM;
        case 7:
          return STREAM_HANDLE;
        case 9:
          return ONE_WAY_MESSAGE;
        case 11:
          return PUSH_DATA;
        case 12:
          return PUSH_MERGED_DATA;
        case 13:
          return REGION_START;
        case 14:
          return REGION_FINISH;
        case 15:
          return PUSH_DATA_HAND_SHAKE;
        case 16:
          return READ_ADD_CREDIT;
        case 17:
          return READ_DATA;
        case 18:
          return OPEN_STREAM_WITH_CREDIT;
        case 19:
          return BACKLOG_ANNOUNCEMENT;
        case 20:
          return TRANSPORTABLE_ERROR;
        case 21:
          return BUFFER_STREAM_END;
        case 22:
          return HEARTBEAT;
        case 23:
          return SEGMENT_START;
        case 24:
          return NOTIFY_REQUIRED_SEGMENT;
        case 25:
          return SUBPARTITION_READ_DATA;
        case -1:
          throw new IllegalArgumentException("User type messages cannot be decoded.");
        default:
          throw new IllegalArgumentException("Unknown message type: " + id);
      }
    }
  }

  public static Message decode(Type msgType, ByteBuf in) {
    return decode(msgType, in, true);
  }

  public static Message decode(Type msgType, ByteBuf in, boolean decodeBody) {
    switch (msgType) {
      case CHUNK_FETCH_REQUEST:
        return ChunkFetchRequest.decode(in);

      case CHUNK_FETCH_SUCCESS:
        return ChunkFetchSuccess.decode(in, decodeBody);

      case CHUNK_FETCH_FAILURE:
        return ChunkFetchFailure.decode(in);

      case RPC_REQUEST:
        return RpcRequest.decode(in, decodeBody);

      case RPC_RESPONSE:
        return RpcResponse.decode(in, decodeBody);

      case RPC_FAILURE:
        return RpcFailure.decode(in);

      case OPEN_STREAM:
        return OpenStream.decode(in);

      case STREAM_HANDLE:
        return StreamHandle.decode(in);

      case ONE_WAY_MESSAGE:
        return OneWayMessage.decode(in, decodeBody);

      case PUSH_DATA:
        return PushData.decode(in, decodeBody);

      case PUSH_MERGED_DATA:
        return PushMergedData.decode(in, decodeBody);

      case REGION_START:
        return RegionStart.decode(in);

      case REGION_FINISH:
        return RegionFinish.decode(in);

      case PUSH_DATA_HAND_SHAKE:
        return PushDataHandShake.decode(in);

      case READ_ADD_CREDIT:
        return ReadAddCredit.decode(in);

      case OPEN_STREAM_WITH_CREDIT:
        return OpenStreamWithCredit.decode(in);

      case BACKLOG_ANNOUNCEMENT:
        return BacklogAnnouncement.decode(in);

      case TRANSPORTABLE_ERROR:
        return TransportableError.decode(in);

      case BUFFER_STREAM_END:
        return BufferStreamEnd.decode(in);

      case HEARTBEAT:
        return new Heartbeat();

      default:
        throw new IllegalArgumentException("Unexpected message type: " + msgType);
    }
  }

  public static Message decode(ByteBuffer buffer) {
    ByteBuf buf = Unpooled.wrappedBuffer(buffer);
    Type type = Type.decode(buf);
    return decode(type, buf);
  }
}
