//
// Copyright Aliaksei Levin (levlam@telegram.org), Arseny Smirnov (arseny30@gmail.com) 2014-2026
//
// Distributed under the Boost Software License, Version 1.0. (See accompanying
// file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)
//
#include "td/e2e/MessageEncryption.h"

#include "EncryptionTestVectors.h"

#include "td/utils/logging.h"
#include "td/utils/misc.h"
#include "td/utils/SharedSlice.h"
#include "td/utils/simple_tests.h"
#include "td/utils/Slice.h"
#include "td/utils/Status.h"
#include "td/utils/tests.h"

namespace tde2e_core {
class EncryptionTest {
 public:
  static td::SecureString encrypt_data_with_deterministic_padding(td::Slice data, td::Slice secret, td::Slice extra) {
    auto prefix = MessageEncryption::gen_deterministic_prefix(data.size(), 16);
    td::SecureString combined(prefix.size() + data.size());
    combined.as_mutable_slice().copy_from(prefix);
    combined.as_mutable_slice().substr(prefix.size()).copy_from(data);
    return MessageEncryption::encrypt_data_with_prefix(combined.as_slice(), secret, extra);
  }
};
}  // namespace tde2e_core

using namespace tde2e_core;

S_TEST(EncryptionTest, test_vectors) {
  auto test_vectors = get_test_vectors();
  for (const auto &vec : test_vectors) {
    LOG(INFO) << "Testing vector: " << vec.name;

    // Convert hex strings to binary
    auto secret = td::hex_decode(vec.secret).move_as_ok();
    auto data = td::hex_decode(vec.data).move_as_ok();
    auto extra = td::hex_decode(vec.extra).move_as_ok();
    auto header = td::hex_decode(vec.header).move_as_ok();
    auto expected_encrypted = td::hex_decode(vec.encrypted).move_as_ok();
    auto expected_encrypted_header = td::hex_decode(vec.encrypted_header).move_as_ok();

    // Test encrypt_data with deterministic padding
    auto encrypted = EncryptionTest::encrypt_data_with_deterministic_padding(data, secret, extra);

    // Test encrypt_header
    auto encrypted_header_result = MessageEncryption::encrypt_header(header, encrypted, secret);
    ASSERT_TRUE(encrypted_header_result.is_ok());

    // For simplicity during debugging, verify only decryption
    auto decrypted_result = MessageEncryption::decrypt_data(expected_encrypted, secret, extra);
    ASSERT_TRUE(decrypted_result.is_ok());
    ASSERT_EQ(td::hex_encode(decrypted_result.ok()), vec.data);

    auto decrypted_header_result =
        MessageEncryption::decrypt_header(expected_encrypted_header, expected_encrypted, secret);
    ASSERT_TRUE(decrypted_header_result.is_ok());
    ASSERT_EQ(td::hex_encode(decrypted_header_result.ok()), vec.header);
  }
  return td::Status::OK();
}