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

#include <string>
#include <utility>
#include <vector>

#include "absl/container/flat_hash_map.h"
#include "absl/functional/any_invocable.h"
#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/cord.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "assertion/assertion_cbor_encoder.h"
#include "assertion/assertion_referencer.h"
#include "async/async_context.h"
#include "claim/signer.h"
#include "constants/labels.h"
#include "jumbf/box_builder.h"
#include "jumbf/constants.h"
#include "jumbf/encode_params.h"
#include "jumbf/utils.h"
#include "proto/actions_assertion.pb.h"
#include "proto/assertion.pb.h"
#include "proto/generator_info.pb.h"
#include "proto/hashed_uri.pb.h"
#include "proto/manifest.pb.h"
#include "proto/multi_asset_hash_assertion.pb.h"
#include "salt/salt_generator.h"
#include "uuid/uuid.h"

namespace credentio {
namespace {

absl::StatusOr<jumbf::BuiltSuperBox> SerializeAssertionStore(
    std::vector<jumbf::BuiltSuperBox> assertions) {
  jumbf::SuperBoxBuilder assertion_store_builder(
      kAssertionStoreUuid, {.label = absl::Cord(kAssertionStoreLabel)});
  for (jumbf::BuiltSuperBox& assertion_box : assertions) {
    ABSL_RETURN_IF_ERROR(
        assertion_store_builder.AddChild(std::move(assertion_box)));
  }
  return std::move(assertion_store_builder).Finalize();
}

absl::StatusOr<jumbf::BuiltSuperBox> SerializeCborSuperBox(
    const Uuid& uuid, absl::string_view label, absl::Cord cbor_content) {
  jumbf::SuperBoxBuilder builder(uuid, {.label = absl::Cord(label)});
  ABSL_RETURN_IF_ERROR(
      builder.AddContent(jumbf::kCborBoxType, std::move(cbor_content)));
  return std::move(builder).Finalize();
}

// Returns description box options for an assertion superbox labeled `label`.
// If `salt_generator` produces a non-empty salt, the salt is embedded as a C2PA
// salt box ('c2sh') in the description box private content, per C2PA
// specification v2.2 sections 8.4.2.3 and 11.1.4.1.2.
absl::StatusOr<jumbf::DescriptionBoxOptions> MakeSaltedDescriptionBoxOptions(
    absl::string_view label, const SaltGenerator& salt_generator) {
  jumbf::DescriptionBoxOptions options = {.label = absl::Cord(label)};
  ABSL_ASSIGN_OR_RETURN(std::string salt, salt_generator.Generate());
  if (salt.empty()) {
    return options;
  }
  ABSL_ASSIGN_OR_RETURN(options.private_content,
                        jumbf::CreateC2paSaltBox(salt));
  return options;
}

// Encodes `assertion` to CBOR with `encoder` and wraps it in a CBOR content
// type JUMBF superbox labeled with the assertion label and salted with
// `salt_generator`.
absl::StatusOr<jumbf::BuiltSuperBox> SerializeAssertionSuperBox(
    const AssertionCborEncoder& encoder, const Assertion& assertion,
    const SaltGenerator& salt_generator) {
  ABSL_ASSIGN_OR_RETURN(absl::Cord cbor, encoder.ToCbor(assertion));
  ABSL_ASSIGN_OR_RETURN(
      jumbf::DescriptionBoxOptions description_options,
      MakeSaltedDescriptionBoxOptions(assertion.label(), salt_generator));
  jumbf::SuperBoxBuilder builder(jumbf::kCborBoxTypeUuid,
                                 std::move(description_options));
  ABSL_RETURN_IF_ERROR(
      builder.AddContent(jumbf::kCborBoxType, std::move(cbor)));
  return std::move(builder).Finalize();
}

struct SignatureBoxResult {
  jumbf::BuiltSuperBox box;
};

void MakeSignatureBox(
    absl::string_view claim_cbor, const ClaimSigner& claim_signer,
    const ClaimSigner::Options& options,
    absl::AnyInvocable<void(absl::StatusOr<SignatureBoxResult>) &&> cb) {
  AsyncResolver<SignatureBoxResult> resolver(options.async_context,
                                             std::move(cb));
  auto done_cb =
      [resolver = std::move(resolver)](
          absl::StatusOr<ClaimSigner::SignResult> signature_status) mutable {
        if (!signature_status.ok()) {
          resolver.Return(signature_status.status());
          return;
        }
        absl::StatusOr<jumbf::BuiltSuperBox> signature_box =
            SerializeCborSuperBox(
                kClaimSignatureUuid, kClaimSignatureLabel,
                absl::Cord(std::move(signature_status->signature)));
        if (!signature_box.ok()) {
          resolver.Return(signature_box.status());
          return;
        }
        resolver.Return(SignatureBoxResult{
            .box = std::move(*signature_box),
        });
      };
  claim_signer.Sign(claim_cbor, options, std::move(done_cb));
}

}  // namespace

absl::Status ManifestBoxEncoder::PopulateMultiAssetPartHashes(
    MultiAssetHashAssertion* assertion) const {
  for (auto& part : *assertion->mutable_parts()) {
    absl::string_view url = part.hash_assertion().url();
    auto it = assertion_url_to_index_.find(url);
    if (it == assertion_url_to_index_.end()) {
      return absl::InternalError(
          absl::StrCat("Could not find hash for url: ", url));
    }
    *part.mutable_hash_assertion() = assertion_hashed_uris_[it->second];
  }
  return absl::OkStatus();
}

absl::Status ManifestBoxEncoder::AddAssertion(Assertion& assertion) {
  if (assertion.has_multi_asset_hash()) {
    ABSL_RETURN_IF_ERROR(
        PopulateMultiAssetPartHashes(assertion.mutable_multi_asset_hash()));
  }

  ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox assertion_box,
                        SerializeAssertionSuperBox(assertion_cbor_encoder_,
                                                   assertion, salt_generator_));
  ABSL_ASSIGN_OR_RETURN(HashedUri assertion_hashed_uri,
                        assertion_reference_.Generate(assertion_box));
  assertion_url_to_index_[assertion_hashed_uri.url()] =
      assertion_hashed_uris_.size();
  assertion_hashed_uris_.Add(std::move(assertion_hashed_uri));
  assertion_boxes_.push_back(std::move(assertion_box));
  return absl::OkStatus();
}

absl::Status ManifestBoxEncoder::AddActionsAssertion(Assertion assertion) {
  for (ActionItem& action : *assertion.mutable_actions()->mutable_actions()) {
    if (action.action() == "c2pa.opened" &&
        action.parameters().ingredients_size() != 1) {
      return absl::InvalidArgumentError(
          "Opened action must have exactly one ingredient.");
    }
    if (action.parameters().ingredients_size() != 0) {
      for (auto& ingredient :
           *action.mutable_parameters()->mutable_ingredients()) {
        auto it = assertion_url_to_index_.find(ingredient.url());
        if (it == assertion_url_to_index_.end()) {
          return absl::InternalError(absl::StrCat(
              "Could not find hash for ingredient url: ", ingredient.url()));
        }
        ingredient = assertion_hashed_uris_[it->second];
      }
    }
    if (action.parameters().has_redacted_uri()) {
      redacted_assertions_.emplace_back(action.parameters().redacted_uri());
    }
  }
  return AddAssertion(assertion);
}

absl::Status ManifestBoxEncoder::AddGatheredAssertion(
    const GatheredAssertion& gathered_assertion) {
  if (gathered_assertion.content_type() !=
      GatheredAssertion::CONTENT_TYPE_JSON) {
    return absl::InvalidArgumentError(
        "Gathered assertion content type must be JSON.");
  }
  ABSL_ASSIGN_OR_RETURN(jumbf::DescriptionBoxOptions description_options,
                        MakeSaltedDescriptionBoxOptions(
                            gathered_assertion.label(), salt_generator_));
  jumbf::SuperBoxBuilder builder(jumbf::kJsonBoxTypeUuid,
                                 std::move(description_options));
  ABSL_RETURN_IF_ERROR(builder.AddContent(
      jumbf::kJsonBoxType, absl::Cord(gathered_assertion.data())));
  ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox assertion_box,
                        std::move(builder).Finalize());
  ABSL_ASSIGN_OR_RETURN(HashedUri assertion_hashed_uri,
                        assertion_reference_.Generate(assertion_box));
  gathered_assertion_hashed_uris_.Add(std::move(assertion_hashed_uri));
  assertion_boxes_.push_back(std::move(assertion_box));
  return absl::OkStatus();
}

absl::Status ManifestBoxEncoder::AddAssertions() {
  // Note that we must add an assertion before adding any other assertion that
  // references it, so the referenced assertion's hash is available. So, for
  // example, ingredient assertions must be added before actions assertions
  // since some actions assertions reference ingredient assertions.
  for (Assertion& assertion : manifest_.ingredient_assertions) {
    ABSL_RETURN_IF_ERROR(AddAssertion(assertion));
  }
  ABSL_RETURN_IF_ERROR(AddActionsAssertion(manifest_.actions_assertion));
  ABSL_RETURN_IF_ERROR(AddAssertion(manifest_.hard_binding_assertion));
  for (Assertion& assertion : manifest_.part_hash_assertions) {
    ABSL_RETURN_IF_ERROR(AddAssertion(assertion));
  }
  if (manifest_.multi_part_assertion.has_value()) {
    ABSL_RETURN_IF_ERROR(AddAssertion(*manifest_.multi_part_assertion));
  }
  for (const auto& gathered_assertion : manifest_.gathered_assertions) {
    ABSL_RETURN_IF_ERROR(AddGatheredAssertion(gathered_assertion));
  }
  return absl::OkStatus();
}

void ManifestBoxEncoder::EncodeClaimAndSignature(
    const ClaimSigner::Options& options,
    absl::AnyInvocable<void(absl::StatusOr<ManifestEncoderResult>) &&> cb) {
  AsyncResolver<ManifestEncoderResult> resolver(options.async_context,
                                                std::move(cb));
  Claim claim;
  claim.set_instance_id(manifest_.instance_id);
  *claim.mutable_claim_generator_info() = generator_info_;
  claim.mutable_created_assertions()->Swap(&assertion_hashed_uris_);
  claim.mutable_gathered_assertions()->Swap(&gathered_assertion_hashed_uris_);
  for (const std::string& redacted_assertion : redacted_assertions_) {
    *claim.add_redacted_assertions() = redacted_assertion;
  }
  claim.set_signature(kClaimSignatureUrl);
  claim.set_default_algorithm(absl::StrCat(hash_algorithm_));

  auto prepare_claim = [&]() -> absl::StatusOr<std::string> {
    ABSL_ASSIGN_OR_RETURN(std::string claim_cbor,
                          claim_cbor_encoder_.ToCbor(claim));
    ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox claim_box,
                          SerializeCborSuperBox(kClaimUuid, kClaimV2Label,
                                                absl::Cord(claim_cbor)));
    ABSL_RETURN_IF_ERROR(manifest_builder_.AddChild(std::move(claim_box)));
    return claim_cbor;
  };
  absl::StatusOr<std::string> claim_cbor = prepare_claim();
  if (!claim_cbor.ok()) {
    resolver.Return(claim_cbor.status());
    return;
  }
  auto done_cb =
      [this, resolver = std::move(resolver)](
          absl::StatusOr<SignatureBoxResult> claim_signature) mutable {
        auto finalize_manifest =
            [&]() -> absl::StatusOr<ManifestEncoderResult> {
          ABSL_RETURN_IF_ERROR(claim_signature.status());
          ABSL_RETURN_IF_ERROR(
              manifest_builder_.AddChild(std::move(claim_signature->box)));
          ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox active_manifest,
                                std::move(manifest_builder_).Finalize());
          ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox finalized_manifest_store,
                                std::move(manifest_.manifest_store_builder)
                                    .Finalize(std::move(active_manifest)));
          return ManifestEncoderResult{
              .box = std::move(finalized_manifest_store),
          };
        };
        resolver.Return(finalize_manifest());
      };
  MakeSignatureBox(*claim_cbor, claim_signer_, options, std::move(done_cb));
}

absl::Status ManifestBoxEncoder::EncodeAssertionStore() {
  ABSL_ASSIGN_OR_RETURN(jumbf::BuiltSuperBox assertion_store,
                        SerializeAssertionStore(std::move(assertion_boxes_)));
  return manifest_builder_.AddChild(std::move(assertion_store));
}

void ManifestBoxEncoder::Encode(
    const ClaimSigner::Options& options,
    absl::AnyInvocable<void(absl::StatusOr<ManifestEncoderResult>) &&> cb) {
  auto add_assertions_status = AddAssertions();
  if (!add_assertions_status.ok()) {
    std::move(cb)(add_assertions_status);
    return;
  }
  auto encode_assertion_store_status = EncodeAssertionStore();
  if (!encode_assertion_store_status.ok()) {
    std::move(cb)(encode_assertion_store_status);
    return;
  }
  EncodeClaimAndSignature(options, std::move(cb));
}
}  // namespace credentio