/*
 * 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 com.datastax.driver.core;

import com.datastax.driver.core.exceptions.DriverInternalError;
import com.datastax.driver.core.exceptions.FrameTooLongException;
import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.handler.codec.ByteToMessageDecoder;
import io.netty.handler.codec.CorruptedFrameException;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import io.netty.handler.codec.MessageToMessageDecoder;
import io.netty.handler.codec.MessageToMessageEncoder;
import io.netty.handler.codec.TooLongFrameException;
import java.util.EnumSet;
import java.util.List;

/**
 * A frame for the CQL binary protocol.
 *
 * <p>Each frame contains a fixed size header (8 bytes for V1 and V2, 9 bytes for V3 and V4)
 * followed by a variable size body. The content of the body depends on the header opcode value (the
 * body can in particular be empty for some opcode values).
 *
 * <p>The protocol distinguishes 2 types of frames: requests and responses. Requests are those
 * frames sent by the clients to the server, response are the ones sent by the server. Note however
 * that the protocol supports server pushes (events) so responses does not necessarily come right
 * after a client request.
 *
 * <p>Frames for protocol versions 1+2 are defined as:
 *
 * <p>
 *
 * <pre>
 *  0         8        16        24        32
 * +---------+---------+---------+---------+
 * | version |  flags  | stream  | opcode  |
 * +---------+---------+---------+---------+
 * |                length                 |
 * +---------+---------+---------+---------+
 * |                                       |
 * .            ...  body ...              .
 * .                                       .
 * .                                       .
 * +---------------------------------------- *
 * </pre>
 *
 * <p>Frames for protocol versions 3+4 are defined as:
 *
 * <p>
 *
 * <pre>
 * 0         8        16        24        32         40
 * +---------+---------+---------+---------+---------+
 * | version |  flags  |      stream       | opcode  |
 * +---------+---------+---------+---------+---------+
 * |                length                 |
 * +---------+---------+---------+---------+
 * |                                       |
 * .            ...  body ...              .
 * .                                       .
 * .                                       .
 * +----------------------------------------
 * </pre>
 *
 * @see "https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v1.spec"
 * @see "https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v2.spec"
 * @see "https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v3.spec"
 * @see "https://github.com/apache/cassandra/blob/trunk/doc/native_protocol_v4.spec"
 */
class Frame {

  final Header header;
  final ByteBuf body;

  Frame(Header header, ByteBuf body) {
    this.header = header;
    this.body = body;
  }

  private static Frame create(ByteBuf fullFrame) {
    Header header = Header.decode(fullFrame);
    assert header.bodyLength == fullFrame.readableBytes();
    return new Frame(header, fullFrame);
  }

  private static int readStreamId(ByteBuf fullFrame, ProtocolVersion version) {
    switch (version) {
      case V1:
      case V2:
        return fullFrame.readByte();
      case V3:
      case V4:
      case V5:
      case V6:
        return fullFrame.readShort();
      default:
        throw version.unsupported();
    }
  }

  static Frame create(
      ProtocolVersion version, int opcode, int streamId, EnumSet<Header.Flag> flags, ByteBuf body) {
    Header header = new Header(version, flags, streamId, opcode, body.readableBytes());
    return new Frame(header, body);
  }

  static class Header {

    final ProtocolVersion version;
    final EnumSet<Flag> flags;
    final int streamId;
    final int opcode;
    final int bodyLength;

    private Header(ProtocolVersion version, int flags, int streamId, int opcode, int bodyLength) {
      this(version, Flag.deserialize(flags), streamId, opcode, bodyLength);
    }

    Header(ProtocolVersion version, EnumSet<Flag> flags, int streamId, int opcode, int bodyLength) {
      this.version = version;
      this.flags = flags;
      this.streamId = streamId;
      this.opcode = opcode;
      this.bodyLength = bodyLength;
    }

    Header withNewBodyLength(int newBodyLength) {
      return new Header(version, flags, streamId, opcode, newBodyLength);
    }

    /**
     * Return the expected frame header length in bytes according to the protocol version in use.
     *
     * @param version the protocol version in use
     * @return the expected frame header length in bytes
     */
    static int lengthFor(ProtocolVersion version) {
      switch (version) {
        case V1:
        case V2:
          return 8;
        case V3:
        case V4:
        case V5:
        case V6:
          return 9;
        default:
          throw version.unsupported();
      }
    }

    public void encodeInto(ByteBuf destination) {
      // Don't bother with the direction, we only send requests.
      destination.writeByte(version.toInt());
      destination.writeByte(Flag.serialize(flags));
      switch (version) {
        case V1:
        case V2:
          destination.writeByte(streamId);
          break;
        case V3:
        case V4:
        case V5:
        case V6:
          destination.writeShort(streamId);
          break;
        default:
          throw version.unsupported();
      }
      destination.writeByte(opcode);
      destination.writeInt(bodyLength);
    }

    static Header decode(ByteBuf buffer) {
      assert buffer.readableBytes() >= 1
          : String.format("Frame too short (%d bytes)", buffer.readableBytes());

      int versionBytes = buffer.readByte();
      // version first byte is the "direction" of the frame (request or response)
      ProtocolVersion version = ProtocolVersion.fromInt(versionBytes & 0x7F);
      int hdrLen = Header.lengthFor(version);
      assert buffer.readableBytes() >= (hdrLen - 1)
          : String.format("Frame too short (%d bytes)", buffer.readableBytes());

      int flags = buffer.readByte();
      int streamId = readStreamId(buffer, version);
      int opcode = buffer.readByte();
      int length = buffer.readInt();

      return new Header(version, flags, streamId, opcode, length);
    }

    enum Flag {
      // The order of that enum matters!!
      COMPRESSED,
      TRACING,
      CUSTOM_PAYLOAD,
      WARNING,
      USE_BETA;

      static EnumSet<Flag> deserialize(int flags) {
        EnumSet<Flag> set = EnumSet.noneOf(Flag.class);
        Flag[] values = Flag.values();
        for (int n = 0; n < 8; n++) {
          if ((flags & (1 << n)) != 0) set.add(values[n]);
        }
        return set;
      }

      static int serialize(EnumSet<Flag> flags) {
        int i = 0;
        for (Flag flag : flags) i |= 1 << flag.ordinal();
        return i;
      }
    }
  }

  Frame with(ByteBuf newBody) {
    return new Frame(header.withNewBodyLength(newBody.readableBytes()), newBody);
  }

  static final class Decoder extends ByteToMessageDecoder {
    private DecoderForStreamIdSize decoder;

    @Override
    protected void decode(ChannelHandlerContext ctx, ByteBuf buffer, List<Object> out)
        throws Exception {
      if (buffer.readableBytes() < 1) return;

      // Initialize sub decoder on first message.  No synchronization needed as
      // decode is always called from same thread.
      if (decoder == null) {
        int version = buffer.getByte(buffer.readerIndex());
        // version first bit is the "direction" of the frame (request or response)
        version = version & 0x7F;
        decoder = new DecoderForStreamIdSize(version, version >= 3 ? 2 : 1);
      }

      Object frame = decoder.decode(ctx, buffer);
      if (frame != null) out.add(frame);
    }

    static class DecoderForStreamIdSize extends LengthFieldBasedFrameDecoder {
      // The maximum response frame length allowed.  Note that C* does not currently restrict the
      // length of its responses (CASSANDRA-12630).
      private static final int MAX_FRAME_LENGTH =
          SystemProperties.getInt("com.datastax.driver.NATIVE_TRANSPORT_MAX_FRAME_SIZE_IN_MB", 256)
              * 1024
              * 1024; // 256 MB
      private final int protocolVersion;

      DecoderForStreamIdSize(int protocolVersion, int streamIdSize) {
        super(MAX_FRAME_LENGTH, /*lengthOffset=*/ 3 + streamIdSize, 4, 0, 0, true);
        this.protocolVersion = protocolVersion;
      }

      @Override
      protected Object decode(ChannelHandlerContext ctx, ByteBuf buffer) throws Exception {
        // Capture current index in case we need to get the stream id.
        // If a TooLongFrameException is thrown the readerIndex will advance to the end of
        // the buffer (or past the frame) so we need the position as we entered this method.
        int curIndex = buffer.readerIndex();
        try {
          ByteBuf frame = (ByteBuf) super.decode(ctx, buffer);
          if (frame == null) {
            return null;
          }
          // Do not deallocate `frame` just yet, because it is stored as Frame.body and will be used
          // in Message.ProtocolDecoder or Frame.Decompressor if compression is enabled (we
          // deallocate
          // it there).
          Frame theFrame = Frame.create(frame);
          // Validate the opcode (this will throw if it's not a response)
          Message.Response.Type.fromOpcode(theFrame.header.opcode);
          return theFrame;
        } catch (CorruptedFrameException e) {
          throw new DriverInternalError(e);
        } catch (TooLongFrameException e) {
          int streamId =
              protocolVersion > 2 ? buffer.getShort(curIndex + 2) : buffer.getByte(curIndex + 2);
          throw new FrameTooLongException(streamId);
        }
      }
    }
  }

  @ChannelHandler.Sharable
  static class Encoder extends MessageToMessageEncoder<Frame> {

    @Override
    protected void encode(ChannelHandlerContext ctx, Frame frame, List<Object> out)
        throws Exception {
      ProtocolVersion protocolVersion = frame.header.version;
      ByteBuf header = ctx.alloc().ioBuffer(Frame.Header.lengthFor(protocolVersion));
      frame.header.encodeInto(header);

      out.add(header);
      out.add(frame.body);
    }
  }

  static class Decompressor extends MessageToMessageDecoder<Frame> {

    private final FrameCompressor compressor;

    Decompressor(FrameCompressor compressor) {
      assert compressor != null;
      this.compressor = compressor;
    }

    @Override
    protected void decode(ChannelHandlerContext ctx, Frame frame, List<Object> out)
        throws Exception {
      if (frame.header.flags.contains(Header.Flag.COMPRESSED)) {
        // All decompressors allocate a new buffer for the decompressed data, so this is the last
        // time
        // we have a reference to the compressed body (and therefore a chance to release it).
        ByteBuf compressedBody = frame.body;
        try {
          out.add(compressor.decompress(frame));
        } finally {
          compressedBody.release();
        }
      } else {
        out.add(frame);
      }
    }
  }

  static class Compressor extends MessageToMessageEncoder<Frame> {

    private final FrameCompressor compressor;

    Compressor(FrameCompressor compressor) {
      assert compressor != null;
      this.compressor = compressor;
    }

    @Override
    protected void encode(ChannelHandlerContext ctx, Frame frame, List<Object> out)
        throws Exception {
      // Never compress STARTUP messages
      if (frame.header.opcode == Message.Request.Type.STARTUP.opcode) {
        out.add(frame);
      } else {
        frame.header.flags.add(Header.Flag.COMPRESSED);
        // See comment in decode()
        ByteBuf uncompressedBody = frame.body;
        try {
          out.add(compressor.compress(frame));
        } finally {
          uncompressedBody.release();
        }
      }
    }
  }
}
