/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * Licensed 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.
 */

// Adapted from Apache Arrow.

#include "velox/dwio/parquet/writer/arrow/FileEncryptorInternal.h"
#include "velox/dwio/parquet/writer/arrow/Encryption.h"
#include "velox/dwio/parquet/writer/arrow/EncryptionInternal.h"

namespace facebook::velox::parquet::arrow {

// Encryptor
Encryptor::Encryptor(
    encryption::AesEncryptor* aes_encryptor,
    const std::string& key,
    const std::string& file_aad,
    const std::string& aad,
    ::arrow::MemoryPool* pool)
    : aes_encryptor_(aes_encryptor),
      key_(key),
      file_aad_(file_aad),
      aad_(aad),
      pool_(pool) {}

int Encryptor::CiphertextSizeDelta() {
  return aes_encryptor_->CiphertextSizeDelta();
}

int Encryptor::Encrypt(
    const uint8_t* plaintext,
    int plaintext_len,
    uint8_t* ciphertext) {
  return aes_encryptor_->Encrypt(
      plaintext,
      plaintext_len,
      str2bytes(key_),
      static_cast<int>(key_.size()),
      str2bytes(aad_),
      static_cast<int>(aad_.size()),
      ciphertext);
}

// InternalFileEncryptor
InternalFileEncryptor::InternalFileEncryptor(
    FileEncryptionProperties* properties,
    ::arrow::MemoryPool* pool)
    : properties_(properties), pool_(pool) {
  if (properties_->is_utilized()) {
    throw ParquetException("Re-using encryption properties for another file");
  }
  properties_->set_utilized();
}

void InternalFileEncryptor::WipeOutEncryptionKeys() {
  properties_->WipeOutEncryptionKeys();

  for (auto const& i : all_encryptors_) {
    i->WipeOut();
  }
}

std::shared_ptr<Encryptor> InternalFileEncryptor::GetFooterEncryptor() {
  if (footer_encryptor_ != nullptr) {
    return footer_encryptor_;
  }

  ParquetCipher::type algorithm = properties_->algorithm().algorithm;
  std::string footer_aad = encryption::CreateFooterAad(properties_->file_aad());
  std::string footer_key = properties_->footer_key();
  auto aes_encryptor = GetMetaAesEncryptor(algorithm, footer_key.size());
  footer_encryptor_ = std::make_shared<Encryptor>(
      aes_encryptor, footer_key, properties_->file_aad(), footer_aad, pool_);
  return footer_encryptor_;
}

std::shared_ptr<Encryptor> InternalFileEncryptor::GetFooterSigningEncryptor() {
  if (footer_signing_encryptor_ != nullptr) {
    return footer_signing_encryptor_;
  }

  ParquetCipher::type algorithm = properties_->algorithm().algorithm;
  std::string footer_aad = encryption::CreateFooterAad(properties_->file_aad());
  std::string footer_signing_key = properties_->footer_key();
  auto aes_encryptor =
      GetMetaAesEncryptor(algorithm, footer_signing_key.size());
  footer_signing_encryptor_ = std::make_shared<Encryptor>(
      aes_encryptor,
      footer_signing_key,
      properties_->file_aad(),
      footer_aad,
      pool_);
  return footer_signing_encryptor_;
}

std::shared_ptr<Encryptor> InternalFileEncryptor::GetColumnMetaEncryptor(
    const std::string& column_path) {
  return GetColumnEncryptor(column_path, true);
}

std::shared_ptr<Encryptor> InternalFileEncryptor::GetColumnDataEncryptor(
    const std::string& column_path) {
  return GetColumnEncryptor(column_path, false);
}

std::shared_ptr<Encryptor>
InternalFileEncryptor::InternalFileEncryptor::GetColumnEncryptor(
    const std::string& column_path,
    bool metadata) {
  // first look if we already got the encryptor from before
  if (metadata) {
    if (column_metadata_map_.find(column_path) != column_metadata_map_.end()) {
      return column_metadata_map_.at(column_path);
    }
  } else {
    if (column_data_map_.find(column_path) != column_data_map_.end()) {
      return column_data_map_.at(column_path);
    }
  }
  auto column_prop = properties_->column_encryption_properties(column_path);
  if (column_prop == nullptr) {
    return nullptr;
  }

  std::string key;
  if (column_prop->is_encrypted_with_footer_key()) {
    key = properties_->footer_key();
  } else {
    key = column_prop->key();
  }

  ParquetCipher::type algorithm = properties_->algorithm().algorithm;
  auto aes_encryptor = metadata ? GetMetaAesEncryptor(algorithm, key.size())
                                : GetDataAesEncryptor(algorithm, key.size());

  std::string file_aad = properties_->file_aad();
  std::shared_ptr<Encryptor> encryptor =
      std::make_shared<Encryptor>(aes_encryptor, key, file_aad, "", pool_);
  if (metadata)
    column_metadata_map_[column_path] = encryptor;
  else
    column_data_map_[column_path] = encryptor;

  return encryptor;
}

int InternalFileEncryptor::MapKeyLenToEncryptorArrayIndex(int key_len) {
  if (key_len == 16)
    return 0;
  else if (key_len == 24)
    return 1;
  else if (key_len == 32)
    return 2;
  throw ParquetException("encryption key must be 16, 24 or 32 bytes in length");
}

encryption::AesEncryptor* InternalFileEncryptor::GetMetaAesEncryptor(
    ParquetCipher::type algorithm,
    size_t key_size) {
  int key_len = static_cast<int>(key_size);
  int index = MapKeyLenToEncryptorArrayIndex(key_len);
  if (meta_encryptor_[index] == nullptr) {
    meta_encryptor_[index].reset(encryption::AesEncryptor::Make(
        algorithm, key_len, true, &all_encryptors_));
  }
  return meta_encryptor_[index].get();
}

encryption::AesEncryptor* InternalFileEncryptor::GetDataAesEncryptor(
    ParquetCipher::type algorithm,
    size_t key_size) {
  int key_len = static_cast<int>(key_size);
  int index = MapKeyLenToEncryptorArrayIndex(key_len);
  if (data_encryptor_[index] == nullptr) {
    data_encryptor_[index].reset(encryption::AesEncryptor::Make(
        algorithm, key_len, false, &all_encryptors_));
  }
  return data_encryptor_[index].get();
}

} // namespace facebook::velox::parquet::arrow
