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

#include <cstdint>
#include <string>
#include <utility>

#include "absl/base/nullability.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "async/async_context.h"
#include "cose/sig_structure.h"
#include "cose/timestamp_creator.h"
#include "cppbor/cppbor.h"

namespace credentio {

void TimestampUpdater::UpdateTimestamp(
    std::string* absl_nonnull encoded_cose_sign1, AsyncContext async_context,
    absl::AnyInvocable<void(absl::Status) &&> cb) const {
  int64_t original_size = encoded_cose_sign1->size();
  auto cose_sign1_or = DecodeCoseSign1TaggedStructure(*encoded_cose_sign1);
  if (!cose_sign1_or.ok()) {
    AsyncResolver<void>(async_context, std::move(cb))
        .Return(cose_sign1_or.status());
    return;
  }
  CoseSign1TaggedStructure cose_sign1 = *std::move(cose_sign1_or);

  absl::AnyInvocable<void(absl::StatusOr<std::string>) &&> done_cb =
      [cose_sign1 = cose_sign1, original_size, encoded_cose_sign1,
       async_context, cb = std::move(cb)](
          absl::StatusOr<std::string> status_or_timestamp) mutable {
        AsyncResolver<void> resolver(async_context, std::move(cb));
        if (!status_or_timestamp.ok()) {
          resolver.Return(status_or_timestamp.status());
          return;
        }
        cose_sign1.unprotected_header = {
            .sig_tst2 =
                TstContainer{
                    .tst_tokens = {{.val = std::move(*status_or_timestamp)}}},
            .pad = "",
        };
        int64_t delta =
            original_size - EncodeCoseSign1TaggedStructure(cose_sign1).size();
        if (delta < 0) {
          resolver.Return(
              absl::InvalidArgumentError("Not enough space for timestamp."));
          return;
        }
        // Adjust padding to account for the encoded size of the `pad` field
        // (size zero is encoded with 1 byte).
        int64_t pad_size = delta - cppbor::Uint(delta).encodedSize() + 1;
        cose_sign1.unprotected_header.pad = std::string(pad_size, '\0');
        *encoded_cose_sign1 = EncodeCoseSign1TaggedStructure(cose_sign1);
        if (encoded_cose_sign1->size() != original_size) {
          resolver.Return(
              absl::InternalError("wrong size after updating timestamp"));
          return;
        }
        resolver.Return(absl::OkStatus());
      };
  timestamp_creator_.CreateTimestamp(
      cose_sign1, TimestampCreator::Options{.async_context = async_context},
      std::move(done_cb));
}

}  // namespace credentio