//===----------------------------------------------------------------------===//

//                         DuckDB

//

// parquet_crypto.hpp

//

//

//===----------------------------------------------------------------------===//

#pragma once

#include "parquet_types.h"

#include "duckdb/common/allocator.hpp"

#include "duckdb/common/encryption_state.hpp"

#include "duckdb/common/encryption_functions.hpp"

#include "duckdb/storage/object_cache.hpp"

namespace duckdb {

class ParquetAdditionalAuthenticatedData;

using duckdb_apache::thrift::TBase;

using duckdb_apache::thrift::protocol::TProtocol;

using duckdb_parquet::ColumnChunk;

using duckdb_parquet::PageType;

class Allocator;

class BufferedFileWriter;

class ParquetKeys : public ObjectCacheEntry {

public:

  static ParquetKeys &Get(ClientContext &context);

public:

  void AddKey(const string &key_name, const string &key);

  bool HasKey(const string &key_name) const;

  const string &GetKey(const string &key_name) const;

public:

  static string ObjectType();

  string GetObjectType() override;

  optional_idx GetEstimatedCacheMemory() const override {

    return optional_idx {};

  }

private:

  unordered_map<string, string> keys;

};

struct CryptoMetaData {

  CryptoMetaData(Allocator &allocator);

  void Initialize(const std::string &unique_file_identifier_p, int16_t row_group_ordinal = -1,

                  int16_t column_ordinal = -1, int8_t module = -1, int16_t page_ordinal = -1);

  void ClearAdditionalAuthenticatedData();

  void SetModule(int8_t module_p);

  bool IsEmpty() const;

public:

  string unique_file_identifier = "";

  int8_t module;

  int16_t row_group_ordinal;

  int16_t column_ordinal;

  int16_t page_ordinal;

public:

  unique_ptr<ParquetAdditionalAuthenticatedData> additional_authenticated_data;

};

class ParquetAdditionalAuthenticatedData : public AdditionalAuthenticatedData {

public:

  explicit ParquetAdditionalAuthenticatedData(Allocator &allocator);

  ~ParquetAdditionalAuthenticatedData() override;

public:

  idx_t GetPrefixSize() const;

  void Rewind() const;

  void WriteParquetAAD(const CryptoMetaData &crypto_meta_data);

private:

  void WritePrefix(const std::string &prefix);

  void WriteSuffix(const CryptoMetaData &crypto_meta_data);

private:

  optional_idx additional_authenticated_data_prefix_size;

};

class ParquetEncryptionConfig {

public:

  explicit ParquetEncryptionConfig();

  ParquetEncryptionConfig(ClientContext &context, const Value &arg);

  ParquetEncryptionConfig(string footer_key);

public:

  static shared_ptr<ParquetEncryptionConfig> Create(ClientContext &context, const Value &arg);

  const string &GetFooterKey() const;

public:

  void Serialize(Serializer &serializer) const;

  static shared_ptr<ParquetEncryptionConfig> Deserialize(Deserializer &deserializer);

private:

  //! The encryption key used for the footer

  string footer_key;

  //! Mapping from column name to key name

  unordered_map<string, string> column_keys;

};

class ParquetCrypto {

public:

  //! Encrypted modules

  static constexpr idx_t LENGTH_BYTES = 4;

  static constexpr idx_t NONCE_BYTES = 12;

  static constexpr idx_t TAG_BYTES = 16;

  //! Block size we encrypt/decrypt

  static constexpr idx_t CRYPTO_BLOCK_SIZE = 4096;

  static constexpr idx_t BLOCK_SIZE = 16;

  // Module types for encryption

  static constexpr int8_t FOOTER = 0;

  static constexpr int8_t COLUMN_METADATA = 1;

  static constexpr int8_t DATA_PAGE = 2;

  static constexpr int8_t DICTIONARY_PAGE = 3;

  static constexpr int8_t DATA_PAGE_HEADER = 4;

  static constexpr int8_t DICTIONARY_PAGE_HEADER = 5;

  static constexpr int8_t COLUMN_INDEX = 6;

  static constexpr int8_t OFFSET_INDEX = 7;

  static constexpr int8_t BLOOM_FILTER_HEADER = 8;

  static constexpr int8_t BLOOM_FILTER_BITSET = 9;

  // Standard AAD length for file

  static constexpr int32_t UNIQUE_FILE_ID_LEN = 8;

  // Maximum Parquet AAD suffix bytes

  static constexpr int32_t AAD_MAX_SUFFIX_BYTES = 7;

public:

  //! Decrypt and read a Thrift object from the transport protocol

  static uint32_t Read(TBase &object, TProtocol &iprot, const string &key, const EncryptionUtil &encryption_util_p,

                       const CryptoMetaData &crypto_meta_data);

  //! Encrypt and write a Thrift object to the transport protocol

  static uint32_t Write(const TBase &object, TProtocol &oprot, const string &key,

                        const EncryptionUtil &encryption_util_p);

  //! Decrypt and read a buffer

  static uint32_t ReadData(TProtocol &iprot, const data_ptr_t buffer, const uint32_t buffer_size, const string &key,

                           const EncryptionUtil &encryption_util_p, const CryptoMetaData &crypto_meta_data);

  //! Encrypt and write a buffer to a file

  static uint32_t WriteData(TProtocol &oprot, const const_data_ptr_t buffer, const uint32_t buffer_size,

                            const string &key, const EncryptionUtil &encryption_util_p);

public:

  static void AddKey(ClientContext &context, const FunctionParameters &parameters);

  static bool ValidKey(const std::string &key);

public:

  static int8_t GetModuleHeader(const ColumnChunk &chunk, uint16_t page_ordinal);

  static int8_t GetModule(const ColumnChunk &chunk, PageType::type page_type, uint16_t page_ordinal);

  static int16_t GetFinalPageOrdinal(const ColumnChunk &chunk, uint8_t module, uint16_t page_ordinal);

  static void GenerateAdditionalAuthenticatedData(Allocator &allocator, CryptoMetaData &aad_crypto_metadata);

  static unique_ptr<ParquetAdditionalAuthenticatedData> GenerateFooterAAD(Allocator &allocator,

                                                                          const std::string &unique_file_identifier);

};

} // namespace duckdb