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

#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 "crypto/hash.h"
#include "google/protobuf/repeated_field.h"
#include "proto/bmff_based_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"

namespace credentio {

namespace {

const int64_t kMaxPowerOfTwo = 0x8000'0000;
const int64_t kMaxBlockCount = 1048576;

absl::StatusOr<int64_t> NextPowerOfTwo(int64_t value) {
  if ((value & (value - 1)) == 0) {
    // value is already a power of two
    return value;
  }
  if (value < 0 || value >= kMaxPowerOfTwo) {
    return absl::InvalidArgumentError(
        absl::StrFormat("Value is negative or too large: %d", value));
  }
  int64_t next_power_of_two = 1;
  while (next_power_of_two < value) {
    next_power_of_two <<= 1;
  }
  return next_power_of_two;
}

absl::Status PadToPowerOfTwo(std::vector<std::string>& entries) {
  ABSL_ASSIGN_OR_RETURN(auto next_power_of_two, NextPowerOfTwo(entries.size()));
  if (next_power_of_two > entries.size()) {
    entries.resize(next_power_of_two, "");
  }
  return absl::OkStatus();
}

absl::StatusOr<std::vector<std::string>> ReduceLayer(
    HasherFactory* factory, std::vector<std::string> layer) {
  std::vector<std::string> reduced_layer;
  for (int64_t i = 0; i + 1 < layer.size(); i += 2) {
    ABSL_ASSIGN_OR_RETURN(auto hasher, factory->Create());
    hasher->Update(layer[i]);
    hasher->Update(layer[i + 1]);
    reduced_layer.push_back(hasher->Digest());
  }
  return reduced_layer;
}

absl::StatusOr<MerkleTree> PopulateLayers(
    HasherFactory* factory, std::vector<std::string> leaf_hashes) {
  MerkleTree tree = {leaf_hashes};
  while (tree[tree.size() - 1].size() > 1) {
    ABSL_ASSIGN_OR_RETURN(auto reduced_layer,
                          ReduceLayer(factory, tree.back()));
    tree.push_back(std::move(reduced_layer));
  }
  return tree;
}

// Returns the number of rows between leaf_row_size and hashes_row_size.
// For example, if leaf_row_size is 16 and hashes_row_size is 4, the function
// will return 4.
absl::StatusOr<int64_t> RowsAway(int64_t leaf_row_size,
                                 int64_t hashes_row_size) {
  if (leaf_row_size <= 0 || hashes_row_size <= 0) {
    return absl::InvalidArgumentError(
        "leaf_row_size and hashes_row_size must be positive");
  }
  if (hashes_row_size > leaf_row_size) {
    return absl::InvalidArgumentError(
        "hashes row size is larger than the leaf row size");
  }
  ABSL_ASSIGN_OR_RETURN(auto leaf_pow2, NextPowerOfTwo(leaf_row_size));
  ABSL_ASSIGN_OR_RETURN(auto hashes_pow2, NextPowerOfTwo(hashes_row_size));
  int64_t rows_away = 0;
  while (leaf_pow2 != hashes_pow2) {
    leaf_pow2 >>= 1;
    rows_away++;
  }
  return rows_away;
}

absl::StatusOr<std::vector<int64_t>> GetFixedBlockSizes(
    uint64_t mdat_box_size, int64_t fixed_block_size) {
  if (fixed_block_size <= 0) {
    return absl::InvalidArgumentError(
        "merkle map has an invalid fixed block size");
  }

  uint64_t num_blocks = mdat_box_size / fixed_block_size;
  if (mdat_box_size % fixed_block_size != 0) {
    num_blocks++;
  }

  if (num_blocks > kMaxBlockCount) {
    return absl::InvalidArgumentError(absl::StrFormat(
        "merkle map block count exceeds maximum allowable limit: %d > %d",
        num_blocks, kMaxBlockCount));
  }

  std::vector<int64_t> block_sizes;
  block_sizes.reserve(num_blocks);
  block_sizes.insert(block_sizes.end(), mdat_box_size / fixed_block_size,
                     fixed_block_size);
  if (mdat_box_size % fixed_block_size != 0) {
    block_sizes.push_back(mdat_box_size % fixed_block_size);
  }
  return block_sizes;
}

absl::StatusOr<std::vector<int64_t>> GetVariableBlockSizes(
    uint64_t mdat_box_size,
    const google::protobuf::RepeatedField<int64_t>& variable_block_sizes) {
  if (variable_block_sizes.size() > kMaxBlockCount) {
    return absl::InvalidArgumentError(absl::StrFormat(
        "merkle map block count exceeds maximum allowable limit: %d > %d",
        variable_block_sizes.size(), kMaxBlockCount));
  }

  int64_t total_block_size = 0;
  std::vector<int64_t> block_sizes;
  block_sizes.reserve(variable_block_sizes.size());
  for (int64_t block_size : variable_block_sizes) {
    if (block_size <= 0) {
      return absl::InvalidArgumentError(
          "merkle map has an invalid variable block size");
    }
    if (total_block_size > INT64_MAX - block_size) {
      return absl::InvalidArgumentError(
          "merkle map's variable block sizes overflow int64_t");
    }
    total_block_size += block_size;
    block_sizes.push_back(block_size);
  }
  if (total_block_size != mdat_box_size) {
    return absl::InvalidArgumentError(
        "merkle map's variable block sizes do not sum to the mdat box size");
  }
  return block_sizes;
}

absl::StatusOr<std::vector<int64_t>> GetSingleBoxSizes(uint64_t mdat_box_size) {
  if (mdat_box_size > INT64_MAX) {
    return absl::InvalidArgumentError(
        "mdat atom's box size is too large to fit in an int64_t");
  }
  return std::vector<int64_t>{static_cast<int64_t>(mdat_box_size)};
}

}  // namespace

absl::StatusOr<std::vector<int64_t>> DeriveMerkleBlockSizes(
    const BmffMerkle& merkle, uint64_t mdat_box_size) {
  if (merkle.has_fixed_block_size() && merkle.variable_block_sizes_size() > 0) {
    return absl::InvalidArgumentError(
        "merkle map has both fixed and variable block sizes");
  }

  std::vector<int64_t> block_sizes;
  if (merkle.has_fixed_block_size()) {
    int64_t fixed_block_size = merkle.fixed_block_size();
    if (fixed_block_size <= 0) {
      return absl::InvalidArgumentError(
          "merkle map has an invalid fixed block size");
    }
    uint64_t num_blocks = mdat_box_size / fixed_block_size;
    if (mdat_box_size % fixed_block_size != 0) {
      num_blocks++;
    }
    if (num_blocks != merkle.count()) {
      return absl::InvalidArgumentError(
          "merkle map's block sizes do not match the expected leaf count");
    }

    ABSL_ASSIGN_OR_RETURN(
        block_sizes,
        GetFixedBlockSizes(mdat_box_size, merkle.fixed_block_size()));
  } else if (merkle.variable_block_sizes_size() > 0) {
    ABSL_ASSIGN_OR_RETURN(
        block_sizes,
        GetVariableBlockSizes(mdat_box_size, merkle.variable_block_sizes()));
  } else {
    ABSL_ASSIGN_OR_RETURN(block_sizes, GetSingleBoxSizes(mdat_box_size));
  }

  if (block_sizes.size() != merkle.count()) {
    return absl::InvalidArgumentError(
        "merkle map's block sizes do not match the expected leaf count");
  }

  return block_sizes;
}

absl::StatusOr<DerivedTreeData> DeriveMerkleTreeData(int64_t leaf_count,
                                                     int64_t hashes_count,
                                                     int64_t auxiliary_count) {
  ABSL_ASSIGN_OR_RETURN(auto leaf_pow2, NextPowerOfTwo(leaf_count));
  ABSL_ASSIGN_OR_RETURN(auto hashes_pow2, NextPowerOfTwo(hashes_count));
  ABSL_ASSIGN_OR_RETURN(auto delta_rows, RowsAway(leaf_pow2, hashes_pow2));
  int64_t expected_hashes_count = leaf_count;
  for (int64_t r = 0; r < delta_rows; ++r) {
    expected_hashes_count = (expected_hashes_count + 1) >> 1;
  }

  if (hashes_count != expected_hashes_count) {
    return absl::InvalidArgumentError(absl::StrFormat(
        "hashes count is not equal to the expected hashes count: "
        "Hashes Count: %d, Leaf Count: %d, Expected Hashes Count: %d",
        hashes_count, leaf_count, expected_hashes_count));
  }

  if ((delta_rows > 0 || auxiliary_count > 0) &&
      auxiliary_count != leaf_count) {
    return absl::InvalidArgumentError(
        absl::StrFormat("incorrect count of auxiliary data boxes: "
                        "Auxiliary Data Count: %d, Leaf Count: %d",
                        auxiliary_count, leaf_count));
  }

  return DerivedTreeData{
      .full_leaf_count = leaf_pow2,
      .full_hashes_count = hashes_pow2,
      .delta_rows = delta_rows,
  };
}

absl::StatusOr<MerkleTree> CreateMerkleTree(HasherFactory* factory,
                                            riegeli::Reader& input,
                                            int64_t starting_offset,
                                            int64_t length,
                                            CreateMerkleTreeOptions options) {
  if (options.fixed_block_size.has_value() &&
      !options.variable_block_sizes.empty()) {
    return absl::InvalidArgumentError(
        "Must specify a fixed block size or variable block sizes, not both.");
  }

  if (!options.fixed_block_size.has_value() &&
      options.variable_block_sizes.empty()) {
    // The entire input is a single leaf node.
    return CreateMerkleTree(factory, input, starting_offset, length,
                            std::vector<int64_t>{length});
  }

  if (options.fixed_block_size.has_value()) {
    if (options.fixed_block_size.value() <= 0) {
      return absl::InvalidArgumentError(
          absl::StrFormat("Fixed block size must be greater than 0. Found: %d",
                          options.fixed_block_size.value()));
    }
    std::vector<int64_t> block_sizes(length / options.fixed_block_size.value(),
                                     options.fixed_block_size.value());
    if (length % options.fixed_block_size.value() != 0) {
      block_sizes.push_back(length % options.fixed_block_size.value());
    }
    return CreateMerkleTree(factory, input, starting_offset, length,
                            block_sizes);
  }

  return CreateMerkleTree(factory, input, starting_offset, length,
                          options.variable_block_sizes);
}

absl::StatusOr<MerkleTree> CreateMerkleTree(HasherFactory* factory,
                                            riegeli::Reader& input,
                                            int64_t starting_offset,
                                            int64_t length,
                                            std::vector<int64_t> block_sizes) {
  int64_t total_blocks_length = 0;
  std::vector<std::string> leaf_hashes;
  leaf_hashes.reserve(block_sizes.size());
  if (!input.Seek(starting_offset) || input.pos() != starting_offset) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to starting offset"));
  }
  for (int64_t block_size : block_sizes) {
    if (block_size <= 0) {
      return absl::InvalidArgumentError(absl::StrFormat(
          "Block size must be greater than 0. Found: %d", block_size));
    }
    total_blocks_length += block_size;
    if (total_blocks_length > length) {
      return absl::InvalidArgumentError(absl::StrFormat(
          "Total block size is too large. Total Block size: %d, Length: %d",
          total_blocks_length, length));
    }
    std::string leaf_data;
    if (!input.Read(block_size, leaf_data)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read leaf data"));
    }
    ABSL_ASSIGN_OR_RETURN(auto hasher, factory->Create());
    hasher->Update(leaf_data);
    leaf_hashes.push_back(hasher->Digest());
  }
  ABSL_RETURN_IF_ERROR(PadToPowerOfTwo(leaf_hashes));

  return PopulateLayers(factory, leaf_hashes);
}

}  // namespace credentio