/**
 * 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.intel.chimera.stream;

import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.nio.channels.ReadableByteChannel;
import java.security.GeneralSecurityException;
import java.security.InvalidAlgorithmParameterException;
import java.security.InvalidKeyException;
import java.util.Properties;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;

import javax.crypto.BadPaddingException;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.ShortBufferException;

import com.google.common.base.Preconditions;
import com.intel.chimera.cipher.Cipher;
import com.intel.chimera.cipher.CipherFactory;
import com.intel.chimera.input.ChannelInput;
import com.intel.chimera.input.Input;
import com.intel.chimera.input.StreamInput;
import com.intel.chimera.utils.Utils;

import static com.intel.chimera.cipher.CipherTransformation.AES_CTR_NOPADDING;

/**
 * PositionedCryptoInputStream provides the capability to decrypt the stream starting
 * at random position as well as provides the foundation for positioned read for
 * decrypting. This needs a stream cipher mode such as AES CTR mode.
 */
public class PositionedCryptoInputStream extends CryptoInputStream {

  /** DirectBuffer pool */
  private final Queue<ByteBuffer> bufferPool = new ConcurrentLinkedQueue<>();

  /** Cipher pool */
  private final Queue<CipherState> cipherPool = new ConcurrentLinkedQueue<>();

  public PositionedCryptoInputStream(Properties props, InputStream in,
      byte[] key, byte[] iv, long streamOffset) throws IOException {
    this(in, Utils.getCipherInstance(AES_CTR_NOPADDING, props),
        Utils.getBufferSize(props), key, iv, streamOffset);
  }

  public PositionedCryptoInputStream(Properties props, ReadableByteChannel in,
      byte[] key, byte[] iv, long streamOffset) throws IOException {
    this(in, Utils.getCipherInstance(AES_CTR_NOPADDING, props),
        Utils.getBufferSize(props), key, iv, streamOffset);
  }

  public PositionedCryptoInputStream(InputStream in, Cipher cipher,
      int bufferSize, byte[] key, byte[] iv, long streamOffset) throws IOException {
    this(new StreamInput(in, bufferSize), cipher, bufferSize, key, iv, streamOffset);
  }

  public PositionedCryptoInputStream(ReadableByteChannel in, Cipher cipher,
      int bufferSize, byte[] key, byte[] iv, long streamOffset) throws IOException {
    this(new ChannelInput(in), cipher, bufferSize, key, iv, streamOffset);
  }

  public PositionedCryptoInputStream(
      Input input,
      Cipher cipher,
      int bufferSize,
      byte[] key,
      byte[] iv,
      long streamOffset) throws IOException {
    super(input, cipher, bufferSize, key, iv, streamOffset);
    Utils.checkStreamCipher(cipher);
  }

  protected long getPos() throws IOException {
    checkStream();
    return streamOffset - outBuffer.remaining();
  }

  /**
   * Read upto the specified number of bytes from a given position
   * within a stream and return the number of bytes read. This does not
   * change the current offset of the stream, and is thread-safe.
   */
  public int read(long position, byte[] buffer, int offset, int length)
      throws IOException {
    checkStream();
    final int n = input.read(position, buffer, offset, length);
    if (n > 0) {
      // This operation does not change the current offset of the file
      decrypt(position, buffer, offset, n);
    }
    return n;
  }

  /**
   * Read the specified number of bytes from a given position within a stream.
   * This does not change the current offset of the stream and is thread-safe.
   */
  public void readFully(long position, byte[] buffer, int offset, int length)
      throws IOException {
    checkStream();
    input.readFully(position, buffer, offset, length);
    if (length > 0) {
      // This operation does not change the current offset of the file
      decrypt(position, buffer, offset, length);
    }
  }

  public void readFully(long position, byte[] buffer) throws IOException {
    readFully(position, buffer, 0, buffer.length);
  }

  public void seek(long position) throws IOException {
    Preconditions.checkArgument(position >= 0, "Cannot seek to negative offset.");
    checkStream();
    /*
     * If data of target pos in the underlying stream has already been read
     * and decrypted in outBuffer, we just need to re-position outBuffer.
     */
    if (position <= getStreamOffset() && position >= getPos()) {
      int forward = (int) (position - getPos());
      if (forward > 0) {
        outBuffer.position(outBuffer.position() + forward);
      }
    } else {
      input.seek(position);
      resetStreamOffset(position);
    }
  }

  /**
   * Decrypt length bytes in buffer starting at offset. Output is also put
   * into buffer starting at offset. It is thread-safe.
   */
  protected void decrypt(long position, byte[] buffer, int offset, int length)
      throws IOException {
    ByteBuffer inBuffer = getBuffer();
    ByteBuffer outBuffer = getBuffer();
    CipherState state = null;
    try {
      state = getCipherState();
      byte[] iv = getInitIV().clone();
      resetCipher(state, position, iv);
      byte padding = getPadding(position);
      inBuffer.position(padding); // Set proper position for input data.

      int n = 0;
      while (n < length) {
        int toDecrypt = Math.min(length - n, inBuffer.remaining());
        inBuffer.put(buffer, offset + n, toDecrypt);

        // Do decryption
        decrypt(state, inBuffer, outBuffer, padding);

        outBuffer.get(buffer, offset + n, toDecrypt);
        n += toDecrypt;
        padding = postDecryption(state, inBuffer, position + n, iv);
      }
    } finally {
      returnBuffer(inBuffer);
      returnBuffer(outBuffer);
      returnCipherState(state);
    }
  }

  /**
   * Do the decryption using inBuffer as input and outBuffer as output.
   * Upon return, inBuffer is cleared; the decrypted data starts at
   * outBuffer.position() and ends at outBuffer.limit();
   */
  private void decrypt(CipherState state, ByteBuffer inBuffer,
      ByteBuffer outBuffer, byte padding) throws IOException {
    Preconditions.checkState(inBuffer.position() >= padding);
    if(inBuffer.position() == padding) {
      // There is no real data in inBuffer.
      return;
    }
    inBuffer.flip();
    outBuffer.clear();
    decryptBuffer(state, inBuffer, outBuffer);
    inBuffer.clear();
    outBuffer.flip();
    if (padding > 0) {
      /*
       * The plain text and cipher text have a 1:1 mapping, they start at the
       * same position.
       */
      outBuffer.position(padding);
    }
  }

  private void decryptBuffer(CipherState state, ByteBuffer inBuffer, ByteBuffer outBuffer)
      throws IOException {
    int inputSize = inBuffer.remaining();
    try {
      int n = state.getCipher().update(inBuffer, outBuffer);
      if (n < inputSize) {
        /**
         * Typically code will not get here. Cipher#update will consume all
         * input data and put result in outBuffer.
         * Cipher#doFinal will reset the cipher context.
         */
        state.getCipher().doFinal(inBuffer, outBuffer);
        state.reset(true);
      }
    } catch (ShortBufferException | IllegalBlockSizeException
        | BadPaddingException e) {
      throw new IOException(e);
    }
  }

  /**
   * This method is executed immediately after decryption. Check whether
   * cipher should be updated and recalculate padding if needed.
   */
  private byte postDecryption(CipherState state, ByteBuffer inBuffer,
      long position, byte[] iv) throws IOException {
    byte padding = 0;
    if (state.isReset()) {
      /*
       * This code is generally not executed since the cipher usually
       * maintains cipher context (e.g. the counter) internally. However,
       * some implementations can't maintain context so a re-init is necessary
       * after each decryption call.
       */
      resetCipher(state, position, iv);
      padding = getPadding(position);
      inBuffer.position(padding);
    }
    return padding;
  }

  /** Calculate the counter and iv, reset the cipher. */
  private void resetCipher(CipherState state, long position, byte[] iv)
      throws IOException {
    final long counter = getCounter(position);
    Utils.calculateIV(getInitIV(), counter, iv);
    try {
      state.getCipher().init(Cipher.DECRYPT_MODE, getKey(), iv);
    } catch (InvalidKeyException | InvalidAlgorithmParameterException e) {
      throw new IOException(e);
    }
    state.reset(false);
  }

  /** Get Cipher from pool */
  private CipherState getCipherState() throws IOException {
    CipherState state = cipherPool.poll();
    if (state == null) {
      Cipher cipher;
      try {
        cipher = CipherFactory.getInstance(getCipher().getTransformation(),
            getCipher().getProperties());
      } catch (GeneralSecurityException e) {
        throw new IOException(e);
      }
      state = new CipherState(cipher);
    }

    return state;
  }

  /** Return Cipher to pool */
  private void returnCipherState(CipherState state) {
    if (state != null) {
      cipherPool.add(state);
    }
  }

  /** Get direct buffer from pool */
  private ByteBuffer getBuffer() {
    ByteBuffer buffer = bufferPool.poll();
    if (buffer == null) {
      buffer = ByteBuffer.allocateDirect(getBufferSize());
    }

    return buffer;
  }

  /** Return direct buffer to pool */
  private void returnBuffer(ByteBuffer buf) {
    if (buf != null) {
      buf.clear();
      bufferPool.add(buf);
    }
  }

  @Override
  public void close() throws IOException {
    if (!isOpen()) {
      return;
    }

    cleanBufferPool();
    super.close();
  }

  /** Clean direct buffer pool */
  private void cleanBufferPool() {
    ByteBuffer buf;
    while ((buf = bufferPool.poll()) != null) {
      Utils.freeDirectBuffer(buf);
    }
  }

  private class CipherState {
    private Cipher cipher;
    private boolean reset;

    public CipherState(Cipher cipher) {
      this.cipher = cipher;
      this.reset = false;
    }

    public Cipher getCipher() {
      return cipher;
    }

    public boolean isReset() {
      return reset;
    }

    public void reset(boolean reset) {
      this.reset = reset;
    }
  }
}