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

#include <cstddef>
#include <cstdint>
#include <optional>
#include <string>

#include "absl/status/status.h"
#include "absl/strings/str_cat.h"
#include "cose/sig_structure.h"
#include "cppbor/cppbor.h"
#include "proto/boxes_hash_assertion.pb.h"
#include "proto/data_hash_assertion.pb.h"
#include "proto/multi_asset_hash_assertion.pb.h"

namespace credentio {
namespace {
// See headerSize() function in cppbor.cpp
constexpr size_t kIntMaxEncodeByteSize = 9;
// The header byte size of a CBOR-encoded `TstContainer` that includes only
// one token of length 0.
constexpr size_t kTstContainerHeaderByteSize = 26;
// This is a value used to pad the empty COSE_Sign1 structure to the expected
// overall size. It is the initial size of pad2 to use when padding a COSE_Sign1
// message and can be decremented to get us to the desired size. This is set to
// 8 because any increase in pad to the expected size can at most result in an
// overshoot of 8 bytes. This occurs when the expected size is 4,294,967,296
// bytes greater than the initial marshalled sign1 size. headerSize(4294967296)
// = 9 and headerSize(0) = 1, so 9 - 1 = 8.
constexpr int kInitialPad2Size = 8;

size_t EncodeByteSize(int64_t v) {
  if (v < 0) {
    return cppbor::Nint(v).encodedSize();
  }
  return cppbor::Uint(v).encodedSize();
}
}  // namespace

void PadDataHashAssertion(DataHashAssertion& data_hash) {
  data_hash.set_pad("");  // https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html
                          // requires its presence.
  data_hash.clear_pad2();
  if (data_hash.exclusions_size() == 0) {
    return;
  }
  size_t pad_size = 0;
  for (const auto& exclusion : data_hash.exclusions()) {
    pad_size += kIntMaxEncodeByteSize - EncodeByteSize(exclusion.start());
    pad_size += kIntMaxEncodeByteSize - EncodeByteSize(exclusion.length());
  }
  data_hash.set_pad(std::string(pad_size, '\0'));
  // `pad` is a Bstr, and its header encodes the length as an Uint, which itself
  // is a variable of encoded sizes. We use an optional `pad2` to count for this
  // variable.
  size_t max_pad_size =
      (kIntMaxEncodeByteSize - 1) * 2 * data_hash.exclusions_size();
  if (cppbor::Uint(max_pad_size).encodedSize() > 1) {
    // The pad_size now becomes a variable of byte sizes, since it can be
    // encoded as 1, 2, 3, 5 or 9 bytes, depending on the value.
    size_t header_delta = cppbor::Uint(max_pad_size).encodedSize() -
                          cppbor::Uint(pad_size).encodedSize();
    data_hash.set_pad2(std::string(header_delta, '\0'));
  }
}

void PadBoxesHashAssertion(BoxesHashAssertion& boxes_hash) {
  for (auto& box : *boxes_hash.mutable_boxes()) {
    // https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html
    // requires `pad` to be set.
    box.set_pad("");
    // Optional, but set it in case some third-party validators require it.
    box.set_pad2("");
  }
}

void PadMultiAssetHashAssertion(MultiAssetHashAssertion& multi_asset_hash) {
  multi_asset_hash.set_pad("");  // https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html
                                 // requires its presence.
  multi_asset_hash.clear_pad2();
  if (multi_asset_hash.parts_size() == 0) {
    return;
  }
  size_t pad_size = 0;
  for (const auto& part : multi_asset_hash.parts()) {
    pad_size +=
        kIntMaxEncodeByteSize - EncodeByteSize(part.location().byte_offset());
    pad_size +=
        kIntMaxEncodeByteSize - EncodeByteSize(part.location().length());
  }
  multi_asset_hash.set_pad(std::string(pad_size, '\0'));
  // `pad` is a Bstr, and its header encodes the length as an Uint, which itself
  // is a variable of encoded sizes. We use an optional `pad2` to count for this
  // variable.
  size_t max_pad_size =
      (kIntMaxEncodeByteSize - 1) * 2 * multi_asset_hash.parts_size();
  if (cppbor::Uint(max_pad_size).encodedSize() > 1) {
    // The pad_size now becomes a variable of byte sizes, since it can be
    // encoded as 1, 2, 3, 5 or 9 bytes, depending on the value.
    size_t header_delta = cppbor::Uint(max_pad_size).encodedSize() -
                          cppbor::Uint(pad_size).encodedSize();
    multi_asset_hash.set_pad2(std::string(header_delta, '\0'));
  }
}

absl::Status PadCoseSign1TaggedStructure(CoseSign1TaggedStructure& cose_sign1,
                                         int64_t max_signature_size,
                                         int64_t max_timestamp_size) {
  if (cose_sign1.unprotected_header.sig_tst2.has_value() &&
      cose_sign1.unprotected_header.sig_tst2->tst_tokens.size() > 1) {
    return absl::InternalError("multiple timestamp tokens");
  }
  if (cose_sign1.signature.size() > max_signature_size) {
    return absl::InternalError(
        "signature size is larger than the maximum allowed.");
  }
  if (cose_sign1.unprotected_header.sig_tst2.has_value() &&
      cose_sign1.unprotected_header.sig_tst2->tst_tokens[0].val.size() >
          max_timestamp_size) {
    return absl::InternalError(
        "timestamp size is larger than the maximum allowed.");
  }

  // Clear paddings.
  cose_sign1.unprotected_header.pad = std::nullopt;
  cose_sign1.unprotected_header.pad2 = std::nullopt;

  // Pad signature.
  int64_t pad_size = max_signature_size - cose_sign1.signature.size();
  // Header size overhead when encoding the signature bstr.
  int64_t header_size = EncodeByteSize(max_signature_size) -
                        EncodeByteSize(cose_sign1.signature.size());

  // Pad timestamp.
  if (cose_sign1.unprotected_header.sig_tst2.has_value()) {
    int64_t timestamp_size =
        cose_sign1.unprotected_header.sig_tst2->tst_tokens[0].val.size();
    pad_size += max_timestamp_size - timestamp_size;
    // Header size overhead when encoding the timestamp bstr.
    header_size +=
        EncodeByteSize(max_timestamp_size) - EncodeByteSize(timestamp_size);
  } else {
    pad_size += max_timestamp_size + kTstContainerHeaderByteSize;
    // Header size overhead when encoding the timestamp bstr.
    header_size += EncodeByteSize(max_timestamp_size) - EncodeByteSize(0);
  }

  // Write padding in the `pad` field.
  cose_sign1.unprotected_header.pad = std::string(pad_size, '\0');

  // Header size overhead when encoding the `pad` bstr.
  int64_t max_pad_size =
      max_signature_size + max_timestamp_size + kTstContainerHeaderByteSize;
  header_size += EncodeByteSize(max_pad_size) - EncodeByteSize(pad_size);

  // Write header overhead padding in the `pad2` field.
  cose_sign1.unprotected_header.pad2 = std::string(header_size, '\0');
  if (EncodeByteSize(header_size) != 1) {
    return absl::InternalError(
        "the additional encoded length of the pad2 header does "
        "not match the assumption of 1.");
  }
  return absl::OkStatus();
}

absl::Status PadCoseSign1TaggedStructure(CoseSign1TaggedStructure& cose_sign1,
                                         size_t target_size) {
  cose_sign1.unprotected_header.pad = "";
  cose_sign1.unprotected_header.pad2 = std::string(kInitialPad2Size, '\0');
  size_t current_size = EncodeCoseSign1TaggedStructure(cose_sign1).size();
  if (target_size < current_size) {
    // This should never actually occur as long as target_size is
    // set to be larger than the 25 bytes (Sign1 structure initial size).
    return absl::InternalError(absl::StrCat("Expected size ", target_size,
                                            " is smaller than base size ",
                                            current_size));
  }
  size_t pad_size = target_size - current_size;
  cose_sign1.unprotected_header.pad = std::string(pad_size, '\0');
  int pad_prefix_length_delta =
      cppbor::Uint(pad_size).encodedSize() - cppbor::Uint(0).encodedSize();
  int new_pad2_size = kInitialPad2Size - pad_prefix_length_delta;
  cose_sign1.unprotected_header.pad2 = std::string(new_pad2_size, '\0');
  return absl::OkStatus();
}

}  // namespace credentio