/*
    * Copyright 2021 The Netty Project
    *
    * The Netty Project 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:
    *
    *   https://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 io.netty.handler.codec.compression;
  
  import com.aayushatharva.brotli4j.encoder.BrotliEncoderChannel;
  import com.aayushatharva.brotli4j.encoder.Encoder;
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.Unpooled;
  import io.netty.channel.Channel;
  import io.netty.channel.ChannelFuture;
  import io.netty.channel.ChannelHandler;
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.channel.ChannelPromise;
  import io.netty.handler.codec.MessageToByteEncoder;
  import io.netty.util.AttributeKey;
  import io.netty.util.ReferenceCountUtil;
  import io.netty.util.internal.ObjectUtil;
  
  import java.io.IOException;
  import java.nio.ByteBuffer;
  import java.nio.channels.ClosedChannelException;
  import java.nio.channels.WritableByteChannel;
  
  /**
   * Compress a {@link ByteBuf} with the Brotli compression.
   * <p>
   * See <a href="https://github.com/google/brotli">brotli</a>.
   */
  @ChannelHandler.Sharable
  public final class BrotliEncoder extends MessageToByteEncoder<ByteBuf> {
  
      private static final AttributeKey<Writer> ATTR = AttributeKey.valueOf("BrotliEncoderWriter");
  
      private final Encoder.Parameters parameters;
      private final boolean isSharable;
      private Writer writer;
  
      /**
       * Create a new {@link BrotliEncoder} Instance with {@link BrotliOptions#DEFAULT}
       * and {@link #isSharable()} set to {@code true}
       */
      public BrotliEncoder() {
          this(BrotliOptions.DEFAULT);
      }
  
      /**
       * Create a new {@link BrotliEncoder} Instance
       *
       * @param brotliOptions {@link BrotliOptions} to use and
       *                      {@link #isSharable()} set to {@code true}
       */
      public BrotliEncoder(BrotliOptions brotliOptions) {
          this(brotliOptions.parameters());
      }
  
      /**
       * Create a new {@link BrotliEncoder} Instance
       * and {@link #isSharable()} set to {@code true}
       *
       * @param parameters {@link Encoder.Parameters} to use
       */
      public BrotliEncoder(Encoder.Parameters parameters) {
          this(parameters, true);
      }
  
      /**
       * <p>
       * Create a new {@link BrotliEncoder} Instance and specify
       * whether this instance will be shared with multiple pipelines or not.
       * </p>
       *
       * If {@link #isSharable()} is true then on {@link #handlerAdded(ChannelHandlerContext)} call,
       * a new {@link Writer} will create, and it will be mapped using {@link Channel#attr(AttributeKey)}
       * so {@link BrotliEncoder} can be shared with multiple pipelines. This works fine but there on every
       * {@link #encode(ChannelHandlerContext, ByteBuf, ByteBuf)} call, we have to get the {@link Writer} associated
       * with the appropriate channel. And this will add a overhead. So it is recommended to set {@link #isSharable()}
       * to {@code false} and create new {@link BrotliEncoder} instance for every pipeline.
       *
       * @param parameters {@link Encoder.Parameters} to use
       * @param isSharable Set to {@code true} if this instance is shared else set to {@code false}
       */
      public BrotliEncoder(Encoder.Parameters parameters, boolean isSharable) {
          super(ByteBuf.class);
          this.parameters = ObjectUtil.checkNotNull(parameters, "Parameters");
          this.isSharable = isSharable;
      }
 
     @Override
     public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
         Writer writer = new Writer(parameters, ctx);
         if (isSharable) {
             ctx.channel().attr(ATTR).set(writer);
         } else {
             this.writer = writer;
         }
         super.handlerAdded(ctx);
     }
 
     @Override
     public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
         finish(ctx);
         super.handlerRemoved(ctx);
     }
 
     @Override
     protected void encode(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out) throws Exception {
         // NO-OP
     }
 
     @Override
     protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect) throws Exception {
         if (!msg.isReadable()) {
             return Unpooled.EMPTY_BUFFER;
         }
 
         Writer writer;
         if (isSharable) {
             writer = ctx.channel().attr(ATTR).get();
         } else {
             writer = this.writer;
         }
 
         // If Writer is 'null' then Writer is not open.
         if (writer == null) {
             return Unpooled.EMPTY_BUFFER;
         } else {
             writer.encode(msg, preferDirect);
             return writer.writableBuffer;
         }
     }
 
     @Override
     public boolean isSharable() {
         return isSharable;
     }
 
     /**
      * Finish the encoding, close streams and write final {@link ByteBuf} to the channel.
      *
      * @param ctx {@link ChannelHandlerContext} which we want to close
      * @throws IOException If an error occurred during closure
      */
     public void finish(ChannelHandlerContext ctx) throws IOException {
         finishEncode(ctx);
     }
 
     private ChannelFuture finishEncode(ChannelHandlerContext ctx) throws IOException {
         Writer writer;
 
         if (isSharable) {
             writer = ctx.channel().attr(ATTR).getAndSet(null);
         } else {
             writer = this.writer;
         }
 
         if (writer != null) {
             writer.close();
             this.writer = null;
             return writer.closeFuture;
         }
         return ctx.newSucceededFuture();
     }
 
     @Override
     public void close(final ChannelHandlerContext ctx, final ChannelPromise promise) throws Exception {
         ChannelFuture f = finishEncode(ctx);
         EncoderUtil.closeAfterFinishEncode(ctx, f, promise);
     }
 
     /**
      * {@link Writer} is the implementation of {@link WritableByteChannel} which encodes
      * Brotli data and stores it into {@link ByteBuf}.
      */
     private static final class Writer implements WritableByteChannel {
 
         private ByteBuf writableBuffer;
         private final BrotliEncoderChannel brotliEncoderChannel;
         private final ChannelHandlerContext ctx;
         private final ChannelPromise closeFuture;
         private boolean closeInitiated;
         private boolean isClosed;
 
         private Writer(Encoder.Parameters parameters, ChannelHandlerContext ctx) throws IOException {
             brotliEncoderChannel = new BrotliEncoderChannel(this, parameters);
             this.ctx = ctx;
             this.closeFuture = ctx.newPromise();
         }
 
         private void encode(ByteBuf msg, boolean preferDirect) throws Exception {
             try {
                 allocate(preferDirect);
 
                 // Compress data and flush it into Buffer.
                 //
                 // As soon as we call flush, Encoder will be triggered to write encoded
                 // data into WritableByteChannel.
                 //
                 // A race condition will not arise because one flush call to encoder will result
                 // in only 1 call at `write(ByteBuffer)`.
                 ByteBuffer nioBuffer = CompressionUtil.safeReadableNioBuffer(msg);
                 int position = nioBuffer.position();
                 brotliEncoderChannel.write(nioBuffer);
                 msg.skipBytes(nioBuffer.position() - position);
                 brotliEncoderChannel.flush();
             } catch (Exception e) {
                 ReferenceCountUtil.release(msg);
                 throw e;
             }
         }
 
         private void allocate(boolean preferDirect) {
             if (preferDirect) {
                 writableBuffer = ctx.alloc().ioBuffer();
             } else {
                 writableBuffer = ctx.alloc().buffer();
             }
         }
 
         @Override
         public int write(ByteBuffer src) throws IOException {
             if (!isOpen()) {
                 throw new ClosedChannelException();
             }
 
             return writableBuffer.writeBytes(src).readableBytes();
         }
 
         @Override
         public boolean isOpen() {
             return !isClosed;
         }
 
         @Override
         public void close() {
             if (closeInitiated) {
                 return;
             }
             closeInitiated = true;
             ctx.executor().execute(new Runnable() {
                 @Override
                 public void run() {
                     try {
                         finish(closeFuture);
                     } catch (IOException ex) {
                         closeFuture.setFailure(new IllegalStateException("Failed to finish encoding", ex));
                     }
                 }
             });
         }
 
         public void finish(final ChannelPromise promise) throws IOException {
             if (!isClosed) {
                 // Allocate a buffer and write last pending data.
                 allocate(true);
 
                 try {
                     brotliEncoderChannel.close();
                     isClosed = true;
                 } catch (Exception ex) {
                     promise.setFailure(ex);
 
                     // Since we have already allocated Buffer for close operation,
                     // we will release that buffer to prevent memory leak.
                     ReferenceCountUtil.release(writableBuffer);
                     return;
                 }
 
                 ctx.writeAndFlush(writableBuffer, promise);
             }
         }
     }
 }