//
//  ========================================================================
//  Copyright (c) 1995-2015 Mort Bay Consulting Pty. Ltd.
//  ------------------------------------------------------------------------
//  All rights reserved. This program and the accompanying materials
//  are made available under the terms of the Eclipse Public License v1.0
//  and Apache License v2.0 which accompanies this distribution.
//
//      The Eclipse Public License is available at
//      http://www.eclipse.org/legal/epl-v10.html
//
//      The Apache License v2.0 is available at
//      http://www.opensource.org/licenses/apache2.0.php
//
//  You may elect to redistribute this code under either of these licenses.
//  ========================================================================
//

package org.eclipse.jetty.websocket.common.extensions.compress;

import java.nio.ByteBuffer;
import java.util.Queue;
import java.util.zip.DataFormatException;
import java.util.zip.Deflater;
import java.util.zip.Inflater;
import java.util.zip.ZipException;

import org.eclipse.jetty.util.BufferUtil;
import org.eclipse.jetty.util.ConcurrentArrayQueue;
import org.eclipse.jetty.util.IteratingCallback;
import org.eclipse.jetty.util.log.Log;
import org.eclipse.jetty.util.log.Logger;
import org.eclipse.jetty.websocket.api.BadPayloadException;
import org.eclipse.jetty.websocket.api.BatchMode;
import org.eclipse.jetty.websocket.api.WriteCallback;
import org.eclipse.jetty.websocket.api.extensions.Frame;
import org.eclipse.jetty.websocket.common.OpCode;
import org.eclipse.jetty.websocket.common.extensions.AbstractExtension;
import org.eclipse.jetty.websocket.common.frames.DataFrame;

public abstract class CompressExtension extends AbstractExtension
{
    protected static final byte[] TAIL_BYTES = new byte[]{0x00, 0x00, (byte)0xFF, (byte)0xFF};
    private static final Logger LOG = Log.getLogger(CompressExtension.class);
    
    /** Never drop tail bytes 0000FFFF, from any frame type */
    protected static final int TAIL_DROP_NEVER = 0;
    /** Always drop tail bytes 0000FFFF, from all frame types */
    protected static final int TAIL_DROP_ALWAYS = 1;
    /** Only drop tail bytes 0000FFFF, from fin==true frames */
    protected static final int TAIL_DROP_FIN_ONLY = 2;

    /** Always set RSV flag, on all frame types */
    protected static final int RSV_USE_ALWAYS = 0;
    /** 
     * Only set RSV flag on first frame in multi-frame messages.
     * <p>
     * Note: this automatically means no-continuation frames have
     * the RSV bit set 
     */
    protected static final int RSV_USE_ONLY_FIRST = 1;

    private final Queue<FrameEntry> entries = new ConcurrentArrayQueue<>();
    private final IteratingCallback flusher = new Flusher();
    private final Deflater compressor;
    private final Inflater decompressor;
    private int tailDrop = TAIL_DROP_NEVER;
    private int rsvUse = RSV_USE_ALWAYS;

    protected CompressExtension()
    {
        compressor = new Deflater(Deflater.BEST_COMPRESSION, true);
        decompressor = new Inflater(true);
        tailDrop = getTailDropMode();
        rsvUse = getRsvUseMode();
    }
    
    public Deflater getDeflater()
    {
        return compressor;
    }

    public Inflater getInflater()
    {
        return decompressor;
    }

    /**
     * Indicates use of RSV1 flag for indicating deflation is in use.
     */
    @Override
    public boolean isRsv1User()
    {
        return true;
    }
    
    /**
     * Return the mode of operation for dropping (or keeping) tail bytes in frames generated by compress (outgoing)
     * 
     * @return either {@link #TAIL_DROP_ALWAYS}, {@link #TAIL_DROP_FIN_ONLY}, or {@link #TAIL_DROP_NEVER}
     */
    abstract int getTailDropMode();

    /**
     * Return the mode of operation for RSV flag use in frames generate by compress (outgoing)
     * 
     * @return either {@link #RSV_USE_ALWAYS} or {@link #RSV_USE_ONLY_FIRST}
     */
    abstract int getRsvUseMode();

    protected void forwardIncoming(Frame frame, ByteAccumulator accumulator)
    {
        DataFrame newFrame = new DataFrame(frame);
        // Unset RSV1 since it's not compressed anymore.
        newFrame.setRsv1(false);

        ByteBuffer buffer = getBufferPool().acquire(accumulator.getLength(), false);
        try
        {
            BufferUtil.flipToFill(buffer);
            accumulator.transferTo(buffer);
            newFrame.setPayload(buffer);
            nextIncomingFrame(newFrame);
        }
        finally
        {
            getBufferPool().release(buffer);
        }
    }

    protected ByteAccumulator decompress(byte[] input)
    {
        // Since we don't track text vs binary vs continuation state, just grab whatever is the greater value.
        int maxSize = Math.max(getPolicy().getMaxTextMessageSize(), getPolicy().getMaxBinaryMessageBufferSize());
        ByteAccumulator accumulator = new ByteAccumulator(maxSize);

        decompressor.setInput(input, 0, input.length);

        if (LOG.isDebugEnabled())
            LOG.debug("Decompressing {} bytes", input.length);

        try
        {
            // It is allowed to send DEFLATE blocks with BFINAL=1.
            // For such blocks, getRemaining() will be > 0 but finished()
            // will be true, so we need to check for both.
            // When BFINAL=0, finished() will always be false and we only
            // check the remaining bytes.
            while (decompressor.getRemaining() > 0 && !decompressor.finished())
            {
                byte[] output = new byte[Math.min(input.length * 2, 32 * 1024)];
                int decompressed = decompressor.inflate(output);
                if (decompressed == 0)
                {
                    if (decompressor.needsInput())
                    {
                        throw new BadPayloadException("Unable to inflate frame, not enough input on frame");
                    }
                    if (decompressor.needsDictionary())
                    {
                        throw new BadPayloadException("Unable to inflate frame, frame erroneously says it needs a dictionary");
                    }
                }
                else
                {
                    accumulator.addChunk(output, 0, decompressed);
                }
            }
            if (LOG.isDebugEnabled())
                LOG.debug("Decompressed {}->{} bytes", input.length, accumulator.getLength());
            return accumulator;
        }
        catch (DataFormatException x)
        {
            throw new BadPayloadException(x);
        }
    }

    @Override
    public void outgoingFrame(Frame frame, WriteCallback callback, BatchMode batchMode)
    {
        // We use a queue and an IteratingCallback to handle concurrency.
        // We must compress and write atomically, otherwise the compression
        // context on the other end gets confused.

        if (flusher.isFailed())
        {
            notifyCallbackFailure(callback, new ZipException());
            return;
        }

        FrameEntry entry = new FrameEntry(frame, callback, batchMode);
        if (LOG.isDebugEnabled())
            LOG.debug("Queuing {}", entry);
        entries.offer(entry);
        flusher.iterate();
    }

    protected void notifyCallbackSuccess(WriteCallback callback)
    {
        try
        {
            if (callback != null)
                callback.writeSuccess();
        }
        catch (Throwable x)
        {
            if (LOG.isDebugEnabled())
                LOG.debug("Exception while notifying success of callback " + callback, x);
        }
    }

    protected void notifyCallbackFailure(WriteCallback callback, Throwable failure)
    {
        try
        {
            if (callback != null)
                callback.writeFailed(failure);
        }
        catch (Throwable x)
        {
            if (LOG.isDebugEnabled())
                LOG.debug("Exception while notifying failure of callback " + callback, x);
        }
    }

    @Override
    public String toString()
    {
        return getClass().getSimpleName();
    }

    private static class FrameEntry
    {
        private final Frame frame;
        private final WriteCallback callback;
        private final BatchMode batchMode;

        private FrameEntry(Frame frame, WriteCallback callback, BatchMode batchMode)
        {
            this.frame = frame;
            this.callback = callback;
            this.batchMode = batchMode;
        }

        @Override
        public String toString()
        {
            return frame.toString();
        }
    }

    private class Flusher extends IteratingCallback implements WriteCallback
    {
        private static final int INPUT_BUFSIZE = 32 * 1024;
        private FrameEntry current;
        private ByteBuffer payload;
        private boolean finished = true;

        @Override
        protected Action process() throws Exception
        {
            if (finished)
            {
                current = entries.poll();
                LOG.debug("Processing {}", current);
                if (current == null)
                    return Action.IDLE;
                deflate(current);
            }
            else
            {
                compress(current, false);
            }
            return Action.SCHEDULED;
        }

        private void deflate(FrameEntry entry)
        {
            Frame frame = entry.frame;
            BatchMode batchMode = entry.batchMode;
            if (OpCode.isControlFrame(frame.getOpCode()) || !frame.hasPayload())
            {
                nextOutgoingFrame(frame, this, batchMode);
                return;
            }

            compress(entry, true);
        }

        private void compress(FrameEntry entry, boolean first)
        {
            // Get a chunk of the payload to avoid to blow
            // the heap if the payload is a huge mapped file.
            Frame frame = entry.frame;
            ByteBuffer data = frame.getPayload();
            int remaining = data.remaining();
            int inputLength = Math.min(remaining, INPUT_BUFSIZE);
            if (LOG.isDebugEnabled())
                LOG.debug("Compressing {}: {} bytes in {} bytes chunk", entry, remaining, inputLength);

            // Avoid to copy the bytes if the ByteBuffer
            // is backed by an array.
            int inputOffset;
            byte[] input;
            if (data.hasArray())
            {
                input = data.array();
                int position = data.position();
                inputOffset = position + data.arrayOffset();
                data.position(position + inputLength);
            }
            else
            {
                input = new byte[inputLength];
                inputOffset = 0;
                data.get(input, 0, inputLength);
            }
            finished = inputLength == remaining;

            compressor.setInput(input, inputOffset, inputLength);

            // Use an additional space in case the content is not compressible.
            byte[] output = new byte[inputLength + 64];
            int outputOffset = 0;
            int outputLength = 0;
            while (true)
            {
                int space = output.length - outputOffset;
                int compressed = compressor.deflate(output, outputOffset, space, Deflater.SYNC_FLUSH);
                outputLength += compressed;
                if (compressed < space)
                {
                    // Everything was compressed.
                    break;
                }
                else
                {
                    // The compressed output is bigger than the uncompressed input.
                    byte[] newOutput = new byte[output.length * 2];
                    System.arraycopy(output, 0, newOutput, 0, output.length);
                    outputOffset += output.length;
                    output = newOutput;
                }
            }

            boolean fin = frame.isFin() && finished;

            // Handle tail bytes generated by SYNC_FLUSH.
            if(tailDrop == TAIL_DROP_ALWAYS) {
                payload = ByteBuffer.wrap(output, 0, outputLength - TAIL_BYTES.length);
            } else if(tailDrop == TAIL_DROP_FIN_ONLY) {
                payload = ByteBuffer.wrap(output, 0, outputLength - (fin?TAIL_BYTES.length:0));
            } else {
                // always include
                payload = ByteBuffer.wrap(output, 0, outputLength);
            }
            if (LOG.isDebugEnabled())
            {
                LOG.debug("Compressed {}: {}->{} chunk bytes",entry,inputLength,outputLength);
            }

            boolean continuation = frame.getType().isContinuation() || !first;
            DataFrame chunk = new DataFrame(frame, continuation);
            if(rsvUse == RSV_USE_ONLY_FIRST) {
                chunk.setRsv1(!continuation);
            } else {
                // always set
                chunk.setRsv1(true);
            }
            chunk.setPayload(payload);
            chunk.setFin(fin);

            nextOutgoingFrame(chunk, this, entry.batchMode);
        }

        @Override
        protected void onCompleteSuccess()
        {
            // This IteratingCallback never completes.
        }
        
        @Override
        protected void onCompleteFailure(Throwable x)
        {
            // Fail all the frames in the queue.
            FrameEntry entry;
            while ((entry = entries.poll()) != null)
                notifyCallbackFailure(entry.callback, x);
        }

        @Override
        public void writeSuccess()
        {
            if (finished)
                notifyCallbackSuccess(current.callback);
            succeeded();
        }

        @Override
        public void writeFailed(Throwable x)
        {
            notifyCallbackFailure(current.callback, x);
            // If something went wrong, very likely the compression context
            // will be invalid, so we need to fail this IteratingCallback.
            failed(x);
        }
    }
}
