/*
 * 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 java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.WritableByteChannel;
import java.util.Locale;

/**
 * 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 BrotliOptions brotliOptions;

    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, true);
    }

    /**
     * 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, 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 isSharable Set to {@code true} if this instance is shared else set to {@code false}
     */
    public BrotliEncoder(BrotliOptions brotliOptions, boolean isSharable) {
        this.brotliOptions = brotliOptions;
        this.isSharable = isSharable;
    }

    @Override
    public void handlerAdded(ChannelHandlerContext ctx) throws Exception {
        Writer writer = new Writer(brotliOptions, 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, ctx.newPromise());
    }

    private ChannelFuture finishEncode(ChannelHandlerContext ctx, ChannelPromise promise) 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 promise;
    }

    @Override
    public void close(final ChannelHandlerContext ctx, final ChannelPromise promise) throws Exception {
        ChannelFuture f = finishEncode(ctx, ctx.newPromise());
        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 boolean isClosed;

        private Writer(BrotliOptions brotliOptions, ChannelHandlerContext ctx) throws IOException {
            Encoder.Parameters parameters = new Encoder.Parameters()
                    .setMode(Encoder.Mode.valueOf(brotliOptions.getMode().toUpperCase(Locale.ROOT)))
                    .setQuality(brotliOptions.getQuality())
                    .setWindow(brotliOptions.getWindow());
            brotliEncoderChannel = new BrotliEncoderChannel(this, parameters);
            this.ctx = ctx;
        }

        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() {
            final ChannelPromise promise = ctx.newPromise();

            ctx.executor().execute(() -> {
                try {
                    finish(promise);
                } catch (IOException ex) {
                    promise.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);
            }
        }
    }
}
