/*
    * 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.decoder.DecoderJNI;
  import io.netty.buffer.ByteBuf;
  import io.netty.channel.ChannelHandlerContext;
  import io.netty.handler.codec.ByteToMessageDecoder;
  import io.netty.util.internal.ObjectUtil;
  
  import java.nio.ByteBuffer;
  import java.util.List;
  
  /**
   * Decompresses a {@link ByteBuf} encoded with the brotli format.
   * <p>
   * See <a href="https://github.com/google/brotli">brotli</a>.
   */
  public final class BrotliDecoder extends ByteToMessageDecoder {
  
      private static final int DEFAULT_MAX_FORWARD_BYTES = CompressionUtil.DEFAULT_MAX_FORWARD_BYTES;
      private static final int DEFAULT_INPUT_BUFFER_SIZE = 8 * 1024;
  
      private enum State {
          DONE, NEEDS_MORE_INPUT, ERROR
      }
  
      static {
          try {
              Brotli.ensureAvailability();
          } catch (Throwable throwable) {
              throw new ExceptionInInitializerError(throwable);
          }
      }
  
      private final int inputBufferSize;
      private final int outputBufferSize;
      private DecoderJNI.Wrapper decoder;
      private boolean destroyed;
      private boolean needsRead;
      private ByteBuf accumBuffer;
  
      /**
       * Creates a new BrotliDecoder with a default 8kB input buffer
       */
      public BrotliDecoder() {
          this(DEFAULT_INPUT_BUFFER_SIZE);
      }
  
      /**
       * Creates a new BrotliDecoder
       * @param inputBufferSize desired size of the input buffer in bytes
       */
      public BrotliDecoder(int inputBufferSize) {
          this(inputBufferSize == 0 ? DEFAULT_INPUT_BUFFER_SIZE : inputBufferSize, DEFAULT_MAX_FORWARD_BYTES);
      }
  
      /**
       * Creates a new BrotliDecoder
       * @param inputBufferSize desired size of the input buffer in bytes
       * @param outputBufferSize desired max size of the output buffer in bytes
       *                         (produce multiple output buffers if exceeded)
       */
      public BrotliDecoder(int inputBufferSize, int outputBufferSize) {
          this.inputBufferSize = ObjectUtil.checkPositive(inputBufferSize, "inputBufferSize");
          this.outputBufferSize = ObjectUtil.checkPositive(outputBufferSize, "outputBufferSize");
      }
  
      /**
       * Creates a new {@link BrotliDecoder} that use the {@code maxAllocation}
       * semantics: the supplied value bounds the size of the emitted decompressed chunks.
       * The input buffer size stays at the decoder's default.
       *
       * @param maxAllocation maximum size, in bytes, of each decompressed output
       *                      buffer forwarded downstream; if {@code 0}, the
       *                      decoder's default output cap is used.
       */
      public static BrotliDecoder newDecoderWithMaxAllocation(int maxAllocation) {
          ObjectUtil.checkPositiveOrZero(maxAllocation, "maxAllocation");
          return maxAllocation > 0 ?
                  new BrotliDecoder(DEFAULT_INPUT_BUFFER_SIZE, maxAllocation) :
                  new BrotliDecoder();
      }
  
      private void forwardOutput(ChannelHandlerContext ctx) {
         ByteBuffer nativeBuffer = decoder.pull(outputBufferSize);
         // nativeBuffer actually wraps brotli's internal buffer so we need to copy its content
         int remaining = nativeBuffer.remaining();
         if (accumBuffer == null) {
             accumBuffer = ctx.alloc().buffer(remaining);
         }
         accumBuffer.writeBytes(nativeBuffer);
         needsRead = false;
         if (accumBuffer.readableBytes() >= outputBufferSize) {
             ctx.fireChannelRead(accumBuffer);
             accumBuffer = null;
         }
     }
 
     private void flushAccumBuffer(ChannelHandlerContext ctx) {
         if (accumBuffer != null && accumBuffer.isReadable()) {
             ctx.fireChannelRead(accumBuffer);
         } else if (accumBuffer != null) {
             accumBuffer.release();
         }
         accumBuffer = null;
     }
 
     private State decompress(ChannelHandlerContext ctx, ByteBuf input) {
         for (;;) {
             switch (decoder.getStatus()) {
                 case DONE:
                     return State.DONE;
 
                 case OK:
                     decoder.push(0);
                     break;
 
                 case NEEDS_MORE_INPUT:
                     while (decoder.hasOutput()) {
                         forwardOutput(ctx);
                     }
 
                     if (!input.isReadable()) {
                         return State.NEEDS_MORE_INPUT;
                     }
 
                     ByteBuffer decoderInputBuffer = decoder.getInputBuffer();
                     decoderInputBuffer.clear();
                     int readBytes = readBytes(input, decoderInputBuffer);
                     decoder.push(readBytes);
                     break;
 
                 case NEEDS_MORE_OUTPUT:
                     forwardOutput(ctx);
                     break;
 
                 default:
                     return State.ERROR;
             }
         }
     }
 
     private static int readBytes(ByteBuf in, ByteBuffer dest) {
         int limit = Math.min(in.readableBytes(), dest.remaining());
         ByteBuffer slice = dest.slice();
         slice.limit(limit);
         in.readBytes(slice);
         dest.position(dest.position() + limit);
         return limit;
     }
 
     @Override
     public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
         decoder = new DecoderJNI.Wrapper(inputBufferSize);
     }
 
     @Override
     protected void decode(ChannelHandlerContext ctx, ByteBuf in, List<Object> out) throws Exception {
         needsRead = true;
         if (destroyed) {
             // Skip data received after finished.
             in.skipBytes(in.readableBytes());
             return;
         }
 
         if (!in.isReadable()) {
             return;
         }
 
         try {
             State state = decompress(ctx, in);
             if (state == State.DONE) {
                 destroy();
             } else if (state == State.ERROR) {
                 throw new DecompressionException("Brotli stream corrupted");
             }
         } catch (Exception e) {
             destroy();
             throw e;
         } finally {
             flushAccumBuffer(ctx);
         }
     }
 
     private void destroy() {
         if (!destroyed) {
             destroyed = true;
             decoder.destroy();
         }
     }
 
     @Override
     protected void handlerRemoved0(ChannelHandlerContext ctx) throws Exception {
         try {
             destroy();
         } finally {
             super.handlerRemoved0(ctx);
         }
     }
 
     @Override
     public void channelInactive(ChannelHandlerContext ctx) throws Exception {
         try {
             destroy();
         } finally {
             super.channelInactive(ctx);
         }
     }
 
     @Override
     public void channelReadComplete(ChannelHandlerContext ctx) throws Exception {
         // Discard bytes of the cumulation buffer if needed.
         discardSomeReadBytes();
 
         if (needsRead && !ctx.channel().config().isAutoRead()) {
             ctx.read();
         }
         ctx.fireChannelReadComplete();
     }
 }