// 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 <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/string_view.h"
#include "bindings/input_hasher.h"
#include "bindings/merkle.h"
#include "cbor/cbor.h"
#include "cbor/parse.h"
#include "constants/labels.h"
#include "constants/status_codes.h"
#include "crypto/hash.h"
#include "formats/bmff/box_header.h"
#include "proto/bmff_based_hash_assertion.cbor.h"
#include "proto/bmff_based_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"
#include "utils/riegeli.h"
#include "utils/status_tracker.h"

namespace credentio {

namespace {

std::string GetMerkleKey(int64_t unique_id, int64_t local_id) {
  return absl::StrFormat("%d-%d", unique_id, local_id);
}

absl::StatusOr<std::string> ReadPurpose(riegeli::Reader& contents) {
  std::string purpose;
  if (!ReadNullTerminatedString(contents, 20, purpose)) {
    return contents.StatusOrAnnotate(absl::DataLossError("kUnexpectedEof"));
  }
  return purpose;
}

absl::StatusOr<BmffMerkleMap> ReadBmffMerkleMap(riegeli::Reader& contents,
                                                int64_t length) {
  constexpr int64_t kMaxMerkleMapSize = 10 * 1024 * 1024;  // 10 MiB
  if (length < 0 || length > kMaxMerkleMapSize) {
    return absl::InvalidArgumentError(
        absl::StrFormat("invalid length for merkle map: %d", length));
  }
  std::string aux_box_raw;
  if (!contents.Read(length, aux_box_raw)) {
    return absl::InternalError("failed to read auxiliary merkle map");
  }

  BmffMerkleMap aux_box;
  ABSL_ASSIGN_OR_RETURN(auto aux_box_cbor, cbor::Parse(aux_box_raw));
  ABSL_ASSIGN_OR_RETURN(auto aux_box_cbor_map, aux_box_cbor->AsMap());
  ABSL_RETURN_IF_ERROR(cbor::ToProto(aux_box_cbor_map, &aux_box));
  return aux_box;
}

absl::StatusOr<std::string> ComputeLeafHash(const HasherFactory& factory,
                                            riegeli::Reader& contents,
                                            int64_t offset, int64_t length) {
  ABSL_ASSIGN_OR_RETURN(auto hasher, factory.Create());
  ABSL_ASSIGN_OR_RETURN(auto input_hasher,
                        InputHasher::Create(std::move(hasher)));
  ABSL_RETURN_IF_ERROR(input_hasher->Update(contents, offset, length));
  return input_hasher->Digest();
}

absl::StatusOr<std::string> JoinHashes(const HasherFactory& factory,
                                       absl::string_view hash1,
                                       absl::string_view hash2) {
  ABSL_ASSIGN_OR_RETURN(auto hasher, factory.Create());
  hasher->Update(hash1);
  hasher->Update(hash2);
  return hasher->Digest();
}

struct LeafData {
  std::vector<BmffBoxHeader> mdat_atoms;
  absl::flat_hash_map<std::string, std::vector<BmffMerkleMap>> auxiliary_data;
};

absl::StatusOr<LeafData> ExtractLeafData(riegeli::Reader& contents) {
  LeafData result;

  if (!contents.Seek(0) || contents.pos() != 0) {
    return absl::InternalError("failed to seek to start of file");
  }
  auto status = IterateOverBmffBoxes(
      contents,
      [&result,
       &contents](const BmffBoxHeader& header) -> absl::StatusOr<bool> {
        // Non-fragmented assets have 1 merkle tree per `mdat`
        // Fragmented assets could have merkle tree per track
        // https://spec.c2pa.org/specifications/specifications/2.2/specs/C2PA_Specification.html#_general_20
        if (header.type == "mdat") {
          // Non-fragmented assets require auxiliary data after the last mdat.
          // https://spec.c2pa.org/specifications/specifications/2.2/specs/C2PA_Specification.html#_non_fragmented_asset_that_can_be_validated_piecewise
          if (!result.auxiliary_data.empty()) {
            return absl::InvalidArgumentError(
                "encountered auxiliary atom before the first mdat atom");
          }
          result.mdat_atoms.push_back(header);
        }
        if (header.type == "uuid" &&
            header.user_type == credentio::kC2paBmffBoxUuid) {
          // Read next string to determine the purpose of the uuid atom.
          ABSL_ASSIGN_OR_RETURN(auto purpose, ReadPurpose(contents));

          if (purpose == "merkle") {
            uint64_t end_of_box = header.start + header.box_size;
            if (contents.pos() >= end_of_box) {
              return absl::InvalidArgumentError(
                  "box size too small for metadata headers");
            }
            uint64_t remaining_data = end_of_box - contents.pos();
            ABSL_ASSIGN_OR_RETURN(auto aux_box,
                                  ReadBmffMerkleMap(contents, remaining_data));

            result
                .auxiliary_data[GetMerkleKey(aux_box.unique_id(),
                                             aux_box.local_id())]
                .push_back(aux_box);
          }
        }
        return true;
      });
  ABSL_RETURN_IF_ERROR(status);
  return result;
}

}  // namespace

absl::Status MerkleValidator::ValidateMerkleMap(
    const BmffMerkle& merkle, const BmffBoxHeader& mdat_atom,
    std::vector<BmffMerkleMap> auxiliary_merkle_maps,
    absl::string_view fallback_algo, StatusTracker& tracker) const {
  if (merkle.has_init_hash()) {
    tracker.RecordFailure(
        FailureStatusCode::kGoogleInternalError,
        {.url = assertion_uri_,
         .explanation = "fragmented merkle validations are not supported yet"});
    return absl::UnimplementedError(
        "fragmented merkle validations are not supported yet");
  }

  // Derive the full row sizes and number of rows away and validates the counts.
  absl::StatusOr<DerivedTreeData> derived_data = DeriveMerkleTreeData(
      merkle.count(), merkle.hashes_size(), auxiliary_merkle_maps.size());
  if (!derived_data.ok()) {
    tracker.RecordFailure(FailureStatusCode::kAssertionBmffHashMalformed,
                          {.url = assertion_uri_,
                           .explanation = derived_data.status().message()});
    return derived_data.status();
  }

  uint64_t mdat_offset_adjustment = 0;
  for (const auto& exclusion : assertion_.exclusions()) {
    if (exclusion.xpath() == "/mdat" && !exclusion.subsets().empty()) {
      mdat_offset_adjustment = exclusion.subsets(0).offset();
      break;
    }
  }

  uint64_t header_or_adjustment_size =
      std::max<uint64_t>(mdat_atom.header_size, mdat_offset_adjustment);

  if (mdat_atom.box_size < header_or_adjustment_size) {
    tracker.RecordFailure(
        FailureStatusCode::kAssertionBmffHashMalformed,
        {.url = assertion_uri_,
         .explanation = "mdat atom box_size is smaller than header_size or "
                        "exclusion adjustment"});
    return absl::InvalidArgumentError(
        "mdat atom box_size is smaller than header_size or exclusion "
        "adjustment");
  }

  // Calculate and validate the lengths of each leaf using net payload size.
  uint64_t net_mdat_payload_size =
      mdat_atom.box_size - header_or_adjustment_size;
  absl::StatusOr<std::vector<int64_t>> leaf_sizes =
      DeriveMerkleBlockSizes(merkle, net_mdat_payload_size);
  if (!leaf_sizes.ok()) {
    tracker.RecordFailure(
        FailureStatusCode::kAssertionBmffHashMalformed,
        {.url = assertion_uri_, .explanation = leaf_sizes.status().message()});
    return leaf_sizes.status();
  }

  // Based on the above calculations, we ensure the assertion is not malformed,
  // now compute and compare the leaf hashes.
  uint64_t offset = mdat_atom.start + header_or_adjustment_size;
  for (int64_t i = 0; i < merkle.count(); ++i) {
    absl::string_view algo = merkle.has_alg() ? merkle.alg() : fallback_algo;
    absl::StatusOr<std::unique_ptr<HasherFactory>> factory =
        factory_provider_->Create(algo);
    if (!factory.ok()) {
      tracker.RecordFailure(FailureStatusCode::kAssertionBmffHashMalformed,
                            {.url = assertion_uri_,
                             .explanation = absl::StrFormat(
                                 "unsupported hash algorithm: %s", algo)});
      return absl::InvalidArgumentError(
          absl::StrFormat("unsupported hash algorithm: %s", algo));
    }

    absl::StatusOr<std::string> leaf_hash =
        ComputeLeafHash(**factory, contents_, offset, leaf_sizes.value()[i]);
    offset += leaf_sizes.value()[i];
    if (!leaf_hash.ok()) {
      tracker.RecordFailure(
          FailureStatusCode::kAssertionBmffHashMalformed,
          {.url = assertion_uri_, .explanation = leaf_hash.status().message()});
      return leaf_hash.status();
    }

    std::string computed_hash = std::move(leaf_hash.value());

    if (!auxiliary_merkle_maps.empty()) {
      const BmffMerkleMap& auxiliary = auxiliary_merkle_maps[i];
      int64_t current_index = auxiliary.location();
      int64_t current_count = merkle.count();
      int proof_index = 0;

      for (int64_t r = 0; r < derived_data->delta_rows; ++r) {
        bool has_sibling = false;
        bool sibling_on_left = false;

        if (current_index % 2 == 1) {
          has_sibling = true;
          sibling_on_left = true;
        } else if (current_index + 1 < current_count) {
          has_sibling = true;
          sibling_on_left = false;
        }

        if (has_sibling) {
          if (proof_index >= auxiliary.hashes_size()) {
            tracker.RecordFailure(
                FailureStatusCode::kAssertionBmffHashMalformed,
                {.url = assertion_uri_,
                 .explanation = "missing auxiliary proof hash"});
            return absl::InvalidArgumentError("missing auxiliary proof hash");
          }
          absl::string_view hash = auxiliary.hashes(proof_index++);
          if (sibling_on_left) {
            ABSL_ASSIGN_OR_RETURN(computed_hash,
                                  JoinHashes(**factory, hash, computed_hash));
          } else {
            ABSL_ASSIGN_OR_RETURN(computed_hash,
                                  JoinHashes(**factory, computed_hash, hash));
          }
        }

        current_index >>= 1;
        current_count = (current_count + 1) >> 1;
      }

      if (proof_index != auxiliary.hashes_size()) {
        tracker.RecordFailure(
            FailureStatusCode::kAssertionBmffHashMalformed,
            {.url = assertion_uri_,
             .explanation = "surplus auxiliary proof hashes"});
        return absl::InvalidArgumentError("surplus auxiliary proof hashes");
      }
    }

    uint64_t hashes_index = i >> derived_data->delta_rows;
    if (merkle.hashes(hashes_index) != computed_hash) {
      tracker.RecordFailure(FailureStatusCode::kAssertionBmffHashMismatch,
                            {.url = assertion_uri_,
                             .explanation = absl::StrFormat(
                                 "merkle map hash mismatch on leaf: %d", i)});
      return absl::InternalError("merkle map hash mismatch");
    }
  }
  return absl::OkStatus();
}

absl::Status MerkleValidator::Validate(StatusTracker& tracker) const {
  absl::StatusOr<LeafData> atoms = ExtractLeafData(contents_);
  if (!atoms.ok()) {
    tracker.RecordFailure(
        FailureStatusCode::kGoogleInternalError,
        {.url = assertion_uri_, .explanation = atoms.status().message()});
    return absl::InternalError("failed to extract leaf data");
  }
  if (atoms->mdat_atoms.size() != assertion_.merkles_size()) {
    // This should only apply to non-fragmented assets.
    tracker.RecordFailure(
        FailureStatusCode::kAssertionBmffHashMalformed,
        {.url = assertion_uri_,
         .explanation =
             "number of mdat atoms does not match the number of merkle maps"});
    return absl::InternalError(
        "number of mdat atoms does not match the number of merkle maps");
  }

  for (int64_t i = 0; i < atoms->mdat_atoms.size(); ++i) {
    BmffMerkle merkle = assertion_.merkles(i);
    const auto& aux_maps = atoms->auxiliary_data[GetMerkleKey(
        merkle.unique_id(), merkle.local_id())];
    ABSL_RETURN_IF_ERROR(ValidateMerkleMap(
        merkle, atoms->mdat_atoms[i], aux_maps, assertion_.alg(), tracker));
  }
  return absl::OkStatus();
}

}  // namespace credentio