// Copyright 2026 Google LLC
//
// 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
//
//    https://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.
//

#include "bindings/merkle.h"

#include <cstdint>
#include <memory>
#include <string>
#include <variant>
#include <vector>

#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "crypto/algorithms.h"
#include "crypto/hash.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/bmff_based_hash_assertion.pb.h"
#include "riegeli/bytes/string_reader.h"
#include "testing/proto_test_utils.h"

namespace credentio {
namespace {

class JoiningHasher : public Hasher {
 public:
  void Update(absl::string_view content) override {
    digest_ = absl::StrCat(digest_, content);
  }

  std::string Digest() override { return digest_; }

 protected:
  std::string digest_ = "";
};

class JoiningHasherFactory : public HasherFactory {
 public:
  absl::StatusOr<std::unique_ptr<Hasher>> Create() const override {
    return std::make_unique<JoiningHasher>();
  }
  HashAlgorithm algorithm() const override { return HashAlgorithm::kSha256; }
};

using ::absl_testing::IsOkAndHolds;
using ::absl_testing::StatusIs;
using ::credentio_testing::ParseTextProtoOrDie;
using ::testing::Eq;
using ::testing::HasSubstr;

struct BlockSizesCase {
  std::string name;
  BmffMerkle merkle_map;
  uint64_t mdat_size;
  absl::StatusOr<std::vector<int64_t>> expected;
};

using DeriveBlockSizesTest = testing::TestWithParam<BlockSizesCase>;

TEST_P(DeriveBlockSizesTest, DerivesBlockSizes) {
  const BlockSizesCase& params = GetParam();

  if (params.expected.ok()) {
    EXPECT_THAT(DeriveMerkleBlockSizes(params.merkle_map, params.mdat_size),
                IsOkAndHolds(Eq(params.expected.value())));
  } else {
    EXPECT_THAT(DeriveMerkleBlockSizes(params.merkle_map, params.mdat_size),
                StatusIs(params.expected.status().code(),
                         params.expected.status().message()));
  }
}

INSTANTIATE_TEST_SUITE_P(
    DeriveMerkleDataTests, DeriveBlockSizesTest,
    ::testing::Values(
        BlockSizesCase{
            .name = "BothSizeTypes",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 5
              fixed_block_size: 10
              variable_block_sizes: 1
              variable_block_sizes: 3
            )pb"),
            .mdat_size = 10,
            .expected = absl::InvalidArgumentError(
                "merkle map has both fixed and variable block sizes"),
        },
        BlockSizesCase{
            .name = "FixedSizeZero",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 5
              fixed_block_size: 0
            )pb"),
            .mdat_size = 10,
            .expected = absl::InvalidArgumentError(
                "merkle map has an invalid fixed block size"),
        },
        BlockSizesCase{
            .name = "FixedSizeBlocksNotMatchingCount",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 3
              fixed_block_size: 3
            )pb"),
            .mdat_size = 10,
            .expected = absl::InvalidArgumentError(
                "merkle map's block sizes do not "
                "match the expected leaf count"),
        },
        BlockSizesCase{
            .name = "VariableHasInvalidSize",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 4
              variable_block_sizes: 3
              variable_block_sizes: 0
              variable_block_sizes: 2
            )pb"),
            .mdat_size = 10,
            .expected = absl::InvalidArgumentError(
                "merkle map has an invalid variable block size"),
        },
        BlockSizesCase{
            .name = "VariableBlockSumTooLarge",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 4
              variable_block_sizes: 3
              variable_block_sizes: 4
              variable_block_sizes: 2
              variable_block_sizes: 6
            )pb"),
            .mdat_size = 10,
            .expected = absl::InvalidArgumentError(
                "merkle map's variable block sizes "
                "do not sum to the mdat box size"),
        },
        BlockSizesCase{
            .name = "ValidVariableBlockSizes",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 4
              variable_block_sizes: 3
              variable_block_sizes: 2
              variable_block_sizes: 3
              variable_block_sizes: 2
            )pb"),
            .mdat_size = 10,
            .expected = std::vector<int64_t>{3, 2, 3, 2},
        },
        BlockSizesCase{
            .name = "NoBlockSizesSpecifiedTooLargeMdat",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 1
            )pb"),
            .mdat_size = static_cast<uint64_t>(INT64_MAX) + 1,
            .expected = absl::InvalidArgumentError(
                "mdat atom's box size is too large to fit in an int64_t"),
        },
        BlockSizesCase{
            .name = "FixedBlockSizeExceedsLimit",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 1048577
              fixed_block_size: 1
            )pb"),
            .mdat_size = 1048577,
            .expected = absl::InvalidArgumentError(
                "merkle map block count exceeds maximum allowable limit: "
                "1048577 > 1048576"),
        },
        BlockSizesCase{
            .name = "VariableBlockSizesOverflow",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 2
              variable_block_sizes: 9223372036854775807
              variable_block_sizes: 1
            )pb"),
            .mdat_size = 10,
            .expected = absl::InvalidArgumentError(
                "merkle map's variable block sizes overflow int64_t"),
        },
        BlockSizesCase{
            .name = "ValidNoBlockSizesSpecified",
            .merkle_map = ParseTextProtoOrDie<BmffMerkle>(R"pb(
              count: 1
            )pb"),
            .mdat_size = 100,
            .expected = std::vector<int64_t>{100},
        }),
    [](const testing::TestParamInfo<BlockSizesCase>& info) {
      return info.param.name;
    });

TEST(MerkleTest, VariableBlockSizesExceedsLimit) {
  BmffMerkle merkle_map;
  merkle_map.set_count(1048577);
  for (int i = 0; i < 1048577; ++i) {
    merkle_map.add_variable_block_sizes(1);
  }
  EXPECT_THAT(DeriveMerkleBlockSizes(merkle_map, 1048577),
              StatusIs(absl::StatusCode::kInvalidArgument,
                       HasSubstr("merkle map block count exceeds maximum "
                                 "allowable limit: 1048577 > 1048576")));
}

struct TreeDataCase {
  std::string name;
  int64_t leaf_count;
  int64_t hashes_count;
  int64_t auxiliary_count;
  absl::StatusOr<DerivedTreeData> expected;
};

using DeriveTreeDataTest = testing::TestWithParam<TreeDataCase>;

TEST_P(DeriveTreeDataTest, DerivesTreeData) {
  const TreeDataCase& params = GetParam();

  if (params.expected.ok()) {
    EXPECT_THAT(DeriveMerkleTreeData(params.leaf_count, params.hashes_count,
                                     params.auxiliary_count),
                IsOkAndHolds(Eq(params.expected.value())));
  } else {
    EXPECT_THAT(DeriveMerkleTreeData(params.leaf_count, params.hashes_count,
                                     params.auxiliary_count),
                StatusIs(params.expected.status().code(),
                         params.expected.status().message()));
  }
}

INSTANTIATE_TEST_SUITE_P(
    DeriveTreeDataTests, DeriveTreeDataTest,
    ::testing::Values(
        TreeDataCase{
            .name = "LeafCountTooLarge",
            .leaf_count = 0x8000'0001,
            .hashes_count = 3,
            .auxiliary_count = 0,
            .expected = absl::InvalidArgumentError(absl::StrFormat(
                "Value is negative or too large: %d", 0x8000'0001)),
        },
        TreeDataCase{
            .name = "HashesCountTooLarge",
            .leaf_count = 5,
            .hashes_count = 0x8000'0001,
            .auxiliary_count = 0,
            .expected = absl::InvalidArgumentError(absl::StrFormat(
                "Value is negative or too large: %d", 0x8000'0001)),
        },
        TreeDataCase{
            .name = "HashesBelowLeafRow",
            .leaf_count = 3,
            .hashes_count = 5,
            .auxiliary_count = 0,
            .expected = absl::InvalidArgumentError(
                "hashes row size is larger than the leaf row size"),
        },
        TreeDataCase{
            .name = "IncorrectHashesCountSameRow",
            .leaf_count = 5,
            .hashes_count = 6,
            .auxiliary_count = 0,
            .expected = absl::InvalidArgumentError(
                "hashes count is not equal to the expected hashes count: "
                "Hashes Count: 6, Leaf Count: 5, Expected Hashes Count: 5"),
        },
        TreeDataCase{
            .name = "IncorrectHashesCountDifferentRows",
            .leaf_count = 5,
            .hashes_count = 4,
            .auxiliary_count = 0,
            .expected = absl::InvalidArgumentError(
                "hashes count is not equal to the expected hashes count: "
                "Hashes Count: 4, Leaf Count: 5, Expected Hashes Count: 3"),
        },
        TreeDataCase{
            .name = "IncorrectAuxiliaryCountSameRows",
            .leaf_count = 5,
            .hashes_count = 5,
            .auxiliary_count = 3,
            .expected = absl::InvalidArgumentError(
                "incorrect count of auxiliary data boxes: Auxiliary Data "
                "Count: 3, Leaf Count: 5"),
        },
        TreeDataCase{
            .name = "IncorrectAuxiliaryCountDifferentRows",
            .leaf_count = 5,
            .hashes_count = 3,
            .auxiliary_count = 3,
            .expected = absl::InvalidArgumentError(
                "incorrect count of auxiliary data boxes: Auxiliary Data "
                "Count: 3, Leaf Count: 5"),
        },
        TreeDataCase{
            .name = "ValidTreeDataSameRowNoAuxiliary",
            .leaf_count = 5,
            .hashes_count = 5,
            .auxiliary_count = 0,
            .expected =
                DerivedTreeData{
                    .full_leaf_count = 8,
                    .full_hashes_count = 8,
                    .delta_rows = 0,
                },
        },
        TreeDataCase{
            .name = "ValidTreeDataSameRowWithAuxiliary",
            .leaf_count = 5,
            .hashes_count = 5,
            .auxiliary_count = 5,
            .expected =
                DerivedTreeData{
                    .full_leaf_count = 8,
                    .full_hashes_count = 8,
                    .delta_rows = 0,
                },
        },
        TreeDataCase{
            .name = "ValidTreeDataDifferentRows",
            .leaf_count = 5,
            .hashes_count = 3,
            .auxiliary_count = 5,
            .expected =
                DerivedTreeData{
                    .full_leaf_count = 8,
                    .full_hashes_count = 4,
                    .delta_rows = 1,
                },
        }),
    [](const testing::TestParamInfo<TreeDataCase>& info) {
      return info.param.name;
    });

struct CreateWithOptionsCase {
  std::string name;
  std::string contents;
  int64_t starting_offset;
  int64_t length;
  CreateMerkleTreeOptions options;
  std::variant<MerkleTree, absl::Status> expected;
};

using CreateWithOptionsTest = testing::TestWithParam<CreateWithOptionsCase>;

TEST_P(CreateWithOptionsTest, CreatesMerkleTree) {
  const CreateWithOptionsCase& params = GetParam();

  riegeli::StringReader<> input(params.contents);

  JoiningHasherFactory factory;

  if (std::holds_alternative<absl::Status>(params.expected)) {
    EXPECT_THAT(
        CreateMerkleTree(&factory, input, params.starting_offset, params.length,
                         params.options),
        StatusIs(std::get<absl::Status>(params.expected).code(),
                 HasSubstr(std::get<absl::Status>(params.expected).message())));
  } else {
    EXPECT_THAT(CreateMerkleTree(&factory, input, params.starting_offset,
                                 params.length, params.options),
                IsOkAndHolds(Eq(std::get<MerkleTree>(params.expected))));
  }
}

INSTANTIATE_TEST_SUITE_P(
    OptionBasedTests, CreateWithOptionsTest,
    ::testing::Values(
        CreateWithOptionsCase{.name = "EmptyOptions",
                              .contents = "abcde",
                              .starting_offset = 0,
                              .length = 5,
                              .options = {},
                              .expected =
                                  MerkleTree{
                                      {"abcde"},
                                  }},
        CreateWithOptionsCase{.name = "FixedBlockSizeSingleByte",
                              .contents = "abcde",
                              .starting_offset = 0,
                              .length = 5,
                              .options = {.fixed_block_size = 1},
                              .expected =
                                  MerkleTree{
                                      {"a", "b", "c", "d", "e", "", "", ""},
                                      {"ab", "cd", "e", ""},
                                      {"abcd", "e"},
                                      {"abcde"},
                                  }},
        CreateWithOptionsCase{.name = "FixedBlockSizeTwoBytes",
                              .contents = "abcde",
                              .starting_offset = 0,
                              .length = 5,
                              .options = {.fixed_block_size = 2},
                              .expected =
                                  MerkleTree{
                                      {"ab", "cd", "e", ""},
                                      {"abcd", "e"},
                                      {"abcde"},
                                  }},
        CreateWithOptionsCase{.name = "VariableBlockSizes",
                              .contents = "abcde",
                              .starting_offset = 0,
                              .length = 5,
                              .options = {.variable_block_sizes = {3, 1, 1}},
                              .expected =
                                  MerkleTree{
                                      {"abc", "d", "e", ""},
                                      {"abcd", "e"},
                                      {"abcde"},
                                  }},
        CreateWithOptionsCase{
            .name = "FailsWithBothFixedAndVariableBlockSizes",
            .contents = "abcde",
            .starting_offset = 0,
            .length = 5,
            .options = {.fixed_block_size = 1,
                        .variable_block_sizes = {3, 1, 1}},
            .expected = absl::InvalidArgumentError(
                "Must specify a fixed block size or "
                "variable block sizes, not both."),
        },
        CreateWithOptionsCase{
            .name = "FailsWithZeroFixedBlockSize",
            .contents = "abcde",
            .starting_offset = 0,
            .length = 5,
            .options = {.fixed_block_size = 0},
            .expected = absl::InvalidArgumentError(
                "Fixed block size must be greater than 0. Found: 0"),
        },
        CreateWithOptionsCase{
            .name = "FailsWithInputTooSmall",
            .contents = "abcde",
            .starting_offset = 3,
            .length = 5,
            .options = {.fixed_block_size = 10},
            .expected = absl::DataLossError("Failed to read leaf data"),
        },
        CreateWithOptionsCase{
            .name = "FixedBlockSizeLargerThanInput",
            .contents = "abcde",
            .starting_offset = 0,
            .length = 5,
            .options = {.fixed_block_size = 10},
            .expected = MerkleTree{{"abcde"}},
        },
        CreateWithOptionsCase{
            .name = "FailsWithVariableBlockSizesTooLarge",
            .contents = "abcde",
            .starting_offset = 0,
            .length = 5,
            .options = {.variable_block_sizes = {2, 3, 5}},
            .expected = absl::InvalidArgumentError(
                "Total block size is too large. "
                "Total Block size: 10, Length: 5"),
        }),
    [](const testing::TestParamInfo<CreateWithOptionsCase>& info) {
      return info.param.name;
    });

}  // namespace
}  // namespace credentio