// 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_validator.h"

#include <algorithm>
#include <cctype>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/status_matchers.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "constants/status_codes.h"
#include "crypto/algorithms.h"
#include "crypto/hash.h"
#include "formats/bmff/test_utils.h"
#include "gmock/gmock.h"
#include "gtest/gtest.h"
#include "proto/bmff_based_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_reader.h"
#include "testing/cbor_utils.h"
#include "testing/proto_test_utils.h"
#include "utils/status_tracker.h"
#include "validator/tracker.h"

namespace credentio {
namespace {

using ::absl_testing::IsOk;
using ::absl_testing::StatusIs;
using ::credentio_testing::ParseTextProtoOrDie;
using ::testing::ElementsAreArray;
using ::testing::IsEmpty;

class TrackingStatusTracker : public StatusTracker {
 public:
  TrackingStatusTracker() = default;
  ~TrackingStatusTracker() override = default;

  struct SuccessRecord {
    SuccessStatusCode code;
    std::string url = "";
    std::string explanation = "";

    bool operator==(const SuccessRecord& other) const {
      return code == other.code && url == other.url &&
             explanation == other.explanation;
    }
  };

  struct FailureRecord {
    FailureStatusCode code;
    std::string url = "";
    std::string explanation = "";

    bool operator==(const FailureRecord& other) const {
      return code == other.code && url == other.url &&
             explanation == other.explanation;
    }
  };

  struct InformationalRecord {
    InformationalStatusCode code;
    std::string url = "";
    std::string explanation = "";

    bool operator==(const InformationalRecord& other) const {
      return code == other.code && url == other.url &&
             explanation == other.explanation;
    }
  };

  void RecordSuccess(SuccessStatusCode code,
                     ValidationTracker::RecordOptions options) override {
    successes_.push_back(SuccessRecord{
        .code = code,
        .url = std::string(options.url),
        .explanation = std::string(options.explanation),
    });
  };
  void RecordFailure(FailureStatusCode code,
                     ValidationTracker::RecordOptions options) override {
    failures_.push_back(FailureRecord{
        .code = code,
        .url = std::string(options.url),
        .explanation = std::string(options.explanation),
    });
  };
  void RecordInformational(InformationalStatusCode code,
                           ValidationTracker::RecordOptions options) override {
    informationals_.push_back(InformationalRecord{
        .code = code,
        .url = std::string(options.url),
        .explanation = std::string(options.explanation),
    });
  };

  std::vector<SuccessRecord> successes_;
  std::vector<FailureRecord> failures_;
  std::vector<InformationalRecord> informationals_;
};

class JoiningHasherFactoryProvider : public HasherFactoryProvider {
 public:
  class JoiningHasherFactory : public HasherFactory {
   public:
    class JoiningHasher : public Hasher {
     public:
      void Update(absl::string_view content) override {
        digest_ = absl::StrCat(digest_, content);

        std::replace_if(
            digest_.begin(), digest_.end(),
            [](char c) { return !std::isalnum(static_cast<unsigned char>(c)); },
            '_');
      }

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

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

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

  absl::StatusOr<std::unique_ptr<HasherFactory>> Create(
      absl::string_view algo) const override {
    if (algo == "error") {
      return absl::InvalidArgumentError(
          absl::StrFormat("unsupported hash algorithm: %s", algo));
    }
    return std::make_unique<JoiningHasherFactory>();
  }
};

struct MerkleValidatorTestCase {
  std::string name;
  std::string file_contents;
  BmffBasedHashAssertion assertion;
  absl::Status expected_status;
  std::vector<TrackingStatusTracker::FailureRecord> expected_failures;
};

class MerkleValidatorTest
    : public ::testing::TestWithParam<MerkleValidatorTestCase> {
 public:
  void SetUp() override {
    auto input_reader =
        std::make_unique<riegeli::StringReader<>>(GetParam().file_contents);
    ASSERT_THAT(input_reader->status(), IsOk());
    input_ = std::move(input_reader);

    factory_ = std::make_unique<JoiningHasherFactoryProvider>();
    assertion_uri_ = "assertion_uri";

    validator_ = std::make_unique<MerkleValidator>(
        *input_, GetParam().assertion, assertion_uri_,
        std::make_unique<JoiningHasherFactoryProvider>());
  }

  std::unique_ptr<riegeli::Reader> input_;
  std::unique_ptr<HasherFactoryProvider> factory_;
  std::string assertion_uri_;
  std::unique_ptr<MerkleValidator> validator_;
};

TEST_P(MerkleValidatorTest, Validate) {
  TrackingStatusTracker tracker;

  EXPECT_THAT(validator_->Validate(tracker),
              StatusIs(GetParam().expected_status.code(),
                       GetParam().expected_status.message()));

  // Successes are not marked within the MerkleValidator.
  EXPECT_THAT(tracker.successes_, IsEmpty());
  EXPECT_THAT(tracker.failures_,
              ElementsAreArray(GetParam().expected_failures));
  EXPECT_THAT(tracker.informationals_, IsEmpty());
}

INSTANTIATE_TEST_SUITE_P(
    MerkleValidatorTests, MerkleValidatorTest,
    ::testing::Values(
        MerkleValidatorTestCase{
            .name = "FailsWithInvalidInput",
            .file_contents = "this_is_not_a_valid_bmff_file",
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles { count: 1 fixed_block_size: 10 }
            )pb"),
            .expected_status =
                absl::InternalError("failed to extract leaf data"),
            .expected_failures =
                {{.code = FailureStatusCode::kGoogleInternalError,
                  .url = "assertion_uri",
                  .explanation = "truncated BMFF box; at byte 8"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsMdatCountNotMatching",
            .file_contents =
                absl::StrCat(credentio_testing::Box("mdat", "data_one"),
                             credentio_testing::Box("mdat", "data_two")),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles { count: 1 fixed_block_size: 10 }
            )pb"),
            .expected_status = absl::InternalError(
                "number of mdat atoms does not match the "
                "number of merkle maps"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMalformed,
                  .url = "assertion_uri",
                  .explanation = "number of mdat atoms does not match "
                                 "the number of merkle maps"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsFragmentedUnsupported",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 2
                init_hash: "only_for_fragmented"
                fixed_block_size: 3
              }
            )pb"),
            .expected_status = absl::UnimplementedError(
                "fragmented merkle validations are not supported yet"),
            .expected_failures =
                {{.code = FailureStatusCode::kGoogleInternalError,
                  .url = "assertion_uri",
                  .explanation =
                      "fragmented merkle validations are not supported yet"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsHashesBelowLeafs",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 3
                fixed_block_size: 10
                hashes: "a"
                hashes: "b"
                hashes: "c"
                hashes: "d"
                hashes: "e"
                hashes: "f"
              }
            )pb"),
            .expected_status = absl::InvalidArgumentError(
                "hashes row size is larger than the leaf row size"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMalformed,
                  .url = "assertion_uri",
                  .explanation =
                      "hashes row size is larger than the leaf row size"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsLeafSizesTooLarge",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles { count: 1 variable_block_sizes: 10 hashes: "abcdef" }
            )pb"),
            .expected_status = absl::InvalidArgumentError(
                "merkle map's variable block sizes do "
                "not sum to the mdat box size"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMalformed,
                  .url = "assertion_uri",
                  .explanation = "merkle map's variable block sizes do not sum "
                                 "to the mdat box size"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsLeafHashMismatch",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 3
                fixed_block_size: 2
                hashes: "wrong_hash"
                hashes: "cd"
                hashes: "ef"
              }
            )pb"),
            .expected_status = absl::InternalError("merkle map hash mismatch"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMismatch,
                  .url = "assertion_uri",
                  .explanation = "merkle map hash mismatch on leaf: 0"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsAuxiliaryDataBeforeLastMdat",
            .file_contents = []() -> std::string {
              std::string aux_box = cbor::FromJson(R"json({
  "hashes": [ "b64'Zm9v'" ],
  "localId": 1,
  "location": 0,
  "uniqueId": 1
})json");
              return absl::StrCat(credentio_testing::Box(
                                      "uuid", credentio_testing::C2paBoxPayload(
                                                  "merkle", aux_box)),
                                  credentio_testing::Box("mdat", "abcdef"));
            }(),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 7
                fixed_block_size: 2
                hashes: "wrong_hash"
                hashes: "cd"
                hashes: "ef"
              }
            )pb"),
            .expected_status =
                absl::InternalError("failed to extract leaf data"),
            .expected_failures =
                {{.code = FailureStatusCode::kGoogleInternalError,
                  .url = "assertion_uri",
                  .explanation =
                      "encountered auxiliary atom before the first mdat atom"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsIncorrectAuxiliaryCount",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 3
                fixed_block_size: 2
                hashes: "wrong_hash"
                hashes: "cd"
                hashes: "ef"
              }
            )pb"),
            .expected_status = absl::InternalError("merkle map hash mismatch"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMismatch,
                  .url = "assertion_uri",
                  .explanation = "merkle map hash mismatch on leaf: 0"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsUnsupportedHashAlgorithm",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "error"
              hash: "unused"
              merkles {
                count: 3
                fixed_block_size: 2
                hashes: "ab"
                hashes: "cd"
                hashes: "ef"
              }
            )pb"),
            .expected_status =
                absl::InvalidArgumentError("unsupported hash algorithm: error"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMalformed,
                  .url = "assertion_uri",
                  .explanation = "unsupported hash algorithm: error"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsInvalidHashWithAuxiliaryData",
            .file_contents = []() -> std::string {
              std::string aux_box_1 = cbor::FromJson(R"json({
  "hashes": [ "b64'ZGVm'" ],
  "localId": 1,
  "location": 0,
  "uniqueId": 1
})json");
              std::string aux_box_2 = cbor::FromJson(R"json({
  "hashes": [ "b64'YWJj'" ],
  "localId": 1,
  "location": 1,
  "uniqueId": 1
})json");
              return absl::StrCat(credentio_testing::Box("mdat", "abcdef"),
                                  credentio_testing::Box(
                                      "uuid", credentio_testing::C2paBoxPayload(
                                                  "merkle", aux_box_1)),
                                  credentio_testing::Box(
                                      "uuid", credentio_testing::C2paBoxPayload(
                                                  "merkle", aux_box_2)));
            }(),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                unique_id: 1
                local_id: 1
                count: 2
                fixed_block_size: 3
                hashes: "wrong_hash"
              }
            )pb"),
            .expected_status = absl::InternalError("merkle map hash mismatch"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMismatch,
                  .url = "assertion_uri",
                  .explanation = "merkle map hash mismatch on leaf: 0"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsBoxSizeTooSmallForMetadataHeaders",
            .file_contents = []() -> std::string {
              std::string aux_box_1 = cbor::FromJson(R"json({
                  "hashes": [ "b64'ZGVm'" ],
                  "localId": 1,
                  "location": 0,
                  "uniqueId": 1
                })json");
              return absl::StrCat(
                  credentio_testing::Box("mdat", "abcdef"),
                  credentio_testing::BadBox(
                      "uuid",
                      credentio_testing::C2paBoxPayload("merkle", aux_box_1),
                      28));
            }(),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                unique_id: 1
                local_id: 1
                count: 2
                fixed_block_size: 3
                hashes: "wrong_hash"
              }
            )pb"),
            .expected_status =
                absl::InternalError("failed to extract leaf data"),
            .expected_failures =
                {{.code = FailureStatusCode::kGoogleInternalError,
                  .url = "assertion_uri",
                  .explanation = "box size too small for metadata headers"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsMdatBoxSizeSmallerThanHeaderSize",
            .file_contents = std::string(
                "\x00\x00\x00\x01mdat\x00\x00\x00\x00\x00\x00\x00\x08", 16),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles { count: 1 fixed_block_size: 2 hashes: "11" }
            )pb"),
            .expected_status =
                absl::InternalError("failed to extract leaf data"),
            .expected_failures =
                {{.code = FailureStatusCode::kGoogleInternalError,
                  .url = "assertion_uri",
                  .explanation = "kInvalidData; box size too small for header; "
                                 "at byte 16"}},
        },
        MerkleValidatorTestCase{
            .name = "FailsTruncatedMerkleCrashPrevention",
            .file_contents = credentio_testing::Box("mdat",
                                                    std::string(28, 'a')),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 14
                fixed_block_size: 2
                hashes: "11"
                hashes: "22"
                hashes: "33"
              }
            )pb"),
            .expected_status = absl::InvalidArgumentError(
                "hashes count is not equal to the expected hashes count: "
                "Hashes Count: 3, Leaf Count: 14, Expected Hashes Count: 4"),
            .expected_failures =
                {{.code = FailureStatusCode::kAssertionBmffHashMalformed,
                  .url = "assertion_uri",
                  .explanation =
                      "hashes count is not equal to the expected hashes count: "
                      "Hashes Count: 3, Leaf Count: 14, Expected Hashes Count: "
                      "4"}},
        },
        MerkleValidatorTestCase{
            .name = "SuccessFixedBlockSizeWithHashesAsLeafRow",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 3
                fixed_block_size: 2
                hashes: "ab"
                hashes: "cd"
                hashes: "ef"
              }
            )pb"),
            .expected_status = absl::OkStatus(),
            .expected_failures = {},
        },
        MerkleValidatorTestCase{
            .name = "SuccessVariableBlockSizeWithHashesAsLeafRow",
            .file_contents = credentio_testing::Box("mdat", "abcdef"),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                count: 2
                variable_block_sizes: 4
                variable_block_sizes: 2
                hashes: "abcd"
                hashes: "ef"
              }
            )pb"),
            .expected_status = absl::OkStatus(),
            .expected_failures = {},
        },
        MerkleValidatorTestCase{
            .name = "SuccessHashWithAuxiliaryData",
            .file_contents = []() -> std::string {
              std::string aux_box_1 = cbor::FromJson(R"json({
  "hashes": [ "b64'ZGVm'" ],
  "localId": 1,
  "location": 0,
  "uniqueId": 1
})json");
              std::string aux_box_2 = cbor::FromJson(R"json({
  "hashes": [ "b64'YWJj'" ],
  "localId": 1,
  "location": 1,
  "uniqueId": 1
})json");
              return absl::StrCat(credentio_testing::Box("mdat", "abcdef"),
                                  credentio_testing::Box(
                                      "uuid", credentio_testing::C2paBoxPayload(
                                                  "merkle", aux_box_1)),
                                  credentio_testing::Box(
                                      "uuid", credentio_testing::C2paBoxPayload(
                                                  "merkle", aux_box_2)));
            }(),
            .assertion = ParseTextProtoOrDie<BmffBasedHashAssertion>(R"pb(
              alg: "sha256"
              hash: "unused"
              merkles {
                unique_id: 1
                local_id: 1
                count: 2
                fixed_block_size: 3
                hashes: "abcdef"
              }
            )pb"),
            .expected_status = absl::OkStatus(),
            .expected_failures = {}}),
    [](const testing::TestParamInfo<MerkleValidatorTestCase>& info) {
      return info.param.name;
    });

using ::testing::NotNull;

TEST(DefaultHasherFactoryProviderTest, CreateSha256) {
  DefaultHasherFactoryProvider provider;
  auto factory = provider.Create("sha256");
  ASSERT_THAT(factory, IsOk());
  EXPECT_THAT(*factory, NotNull());
  EXPECT_EQ((*factory)->algorithm(), HashAlgorithm::kSha256);
}

TEST(DefaultHasherFactoryProviderTest, CreateSha384) {
  DefaultHasherFactoryProvider provider;
  auto factory = provider.Create("sha384");
  ASSERT_THAT(factory, IsOk());
  EXPECT_THAT(*factory, NotNull());
  EXPECT_EQ((*factory)->algorithm(), HashAlgorithm::kSha384);
}

TEST(DefaultHasherFactoryProviderTest, CreateSha512) {
  DefaultHasherFactoryProvider provider;
  auto factory = provider.Create("sha512");
  ASSERT_THAT(factory, IsOk());
  EXPECT_THAT(*factory, NotNull());
  EXPECT_EQ((*factory)->algorithm(), HashAlgorithm::kSha512);
}

TEST(DefaultHasherFactoryProviderTest, CreateInvalid) {
  DefaultHasherFactoryProvider provider;
  EXPECT_THAT(provider.Create("invalid"),
              StatusIs(absl::StatusCode::kInvalidArgument));
}

}  // namespace
}  // namespace credentio