// 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 "generator/generator_impl.h"

#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>

#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/string_view.h"
#include "absl/strings/substitute.h"
#include "async/async_context.h"
#include "cbor/options.h"
#include "claim/manifest_box.h"
#include "claim/manifest_box_encoder.h"
#include "claim/manifest_store_view.h"
#include "cose/sig_structure.h"
#include "generator/manifest_store_transformer.h"
#include "proto/bmff_based_hash_assertion.cbor.h"
#include "proto/bmff_based_hash_assertion.pb.h"
#include "proto/boxes_hash_assertion.cbor.h"
#include "proto/boxes_hash_assertion.pb.h"
#include "proto/collection_data_hash_assertion.cbor.h"
#include "proto/collection_data_hash_assertion.pb.h"
#include "proto/data_hash_assertion.cbor.h"
#include "proto/data_hash_assertion.pb.h"
#include "proto/generator.pb.h"
#include "utils/cbor_padding.h"

namespace credentio {

namespace {

// Current default limit on manifest store size in our validator is 10 MiB.
// Some PNG parsing code elsewhere fails if the manifest store is > 7.6 MiB.
constexpr int64_t kDefaultMaxManifestStoreSizeBytes = 7 * 1024 * 1024;  // 7 MiB

template <typename T>
absl::StatusOr<size_t> GetCborSize(const T& proto,
                                   const cbor::FromProtoOptions& options) {
  ABSL_ASSIGN_OR_RETURN(auto cbor, cbor::FromProto(proto, options));
  return cbor->toString().size();
}

absl::Status CheckIfPartBindingByteSizesMatched(
    const PartBinding& template_part_binding,
    const PartBinding& finalize_part_binding,
    const cbor::FromProtoOptions& proto_to_cbor_options) {
  size_t template_size;
  size_t finalize_size;
  switch (template_part_binding.binding_case()) {
    case PartBinding::kBmffBasedHash: {
      ABSL_ASSIGN_OR_RETURN(template_size,
                            GetCborSize(template_part_binding.bmff_based_hash(),
                                        proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(finalize_size,
                            GetCborSize(finalize_part_binding.bmff_based_hash(),
                                        proto_to_cbor_options));
      break;
    }
    case PartBinding::kBoxHash: {
      BoxesHashAssertion template_boxes_hash = template_part_binding.box_hash();
      BoxesHashAssertion finalize_boxes_hash = finalize_part_binding.box_hash();
      PadBoxesHashAssertion(template_boxes_hash);
      PadBoxesHashAssertion(finalize_boxes_hash);
      ABSL_ASSIGN_OR_RETURN(template_size, GetCborSize(template_boxes_hash,
                                                       proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(finalize_size, GetCborSize(finalize_boxes_hash,
                                                       proto_to_cbor_options));
      break;
    }
    case PartBinding::kDataHash: {
      DataHashAssertion template_data_hash = template_part_binding.data_hash();
      DataHashAssertion finalize_data_hash = finalize_part_binding.data_hash();
      PadDataHashAssertion(template_data_hash);
      PadDataHashAssertion(finalize_data_hash);
      ABSL_ASSIGN_OR_RETURN(template_size, GetCborSize(template_data_hash,
                                                       proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(finalize_size, GetCborSize(finalize_data_hash,
                                                       proto_to_cbor_options));
      break;
    }
    case PartBinding::BINDING_NOT_SET:
      return absl::InvalidArgumentError(
          "A part binding in the finalize request must be set");
  }
  if (template_size != finalize_size) {
    return absl::InvalidArgumentError(
        "The part binding differs from template in unsupported ways.");
  }
  return absl::OkStatus();
}

absl::Status CheckIfHardBindingByteSizesMatched(
    const HardBinding& template_hard_binding,
    const HardBinding& finalize_hard_binding,
    const cbor::FromProtoOptions& proto_to_cbor_options) {
  size_t template_size;
  size_t finalize_size;
  switch (template_hard_binding.binding_case()) {
    case HardBinding::kBmffBasedHash: {
      ABSL_ASSIGN_OR_RETURN(template_size,
                            GetCborSize(template_hard_binding.bmff_based_hash(),
                                        proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(finalize_size,
                            GetCborSize(finalize_hard_binding.bmff_based_hash(),
                                        proto_to_cbor_options));
      break;
    }
    case HardBinding::kBoxHash: {
      BoxesHashAssertion template_boxes_hash = template_hard_binding.box_hash();
      BoxesHashAssertion finalize_boxes_hash = finalize_hard_binding.box_hash();
      PadBoxesHashAssertion(template_boxes_hash);
      PadBoxesHashAssertion(finalize_boxes_hash);
      ABSL_ASSIGN_OR_RETURN(template_size, GetCborSize(template_boxes_hash,
                                                       proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(finalize_size, GetCborSize(finalize_boxes_hash,
                                                       proto_to_cbor_options));
      break;
    }
    case HardBinding::kDataHash: {
      DataHashAssertion template_data_hash = template_hard_binding.data_hash();
      DataHashAssertion finalize_data_hash = finalize_hard_binding.data_hash();
      PadDataHashAssertion(template_data_hash);
      PadDataHashAssertion(finalize_data_hash);
      ABSL_ASSIGN_OR_RETURN(template_size, GetCborSize(template_data_hash,
                                                       proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(finalize_size, GetCborSize(finalize_data_hash,
                                                       proto_to_cbor_options));
      break;
    }
    case HardBinding::kCollectionDataHash: {
      ABSL_ASSIGN_OR_RETURN(
          template_size,
          GetCborSize(template_hard_binding.collection_data_hash(),
                      proto_to_cbor_options));
      ABSL_ASSIGN_OR_RETURN(
          finalize_size,
          GetCborSize(finalize_hard_binding.collection_data_hash(),
                      proto_to_cbor_options));
      break;
    }
    case HardBinding::BINDING_NOT_SET:
      return absl::InvalidArgumentError(
          "The hard binding in the finalize request must be set");
  }
  if (template_size != finalize_size) {
    return absl::InvalidArgumentError(
        "The final hard binding differs from template in unsupported ways.");
  }

  if (template_hard_binding.part_bindings_size() !=
      finalize_hard_binding.part_bindings_size()) {
    return absl::InvalidArgumentError(
        "The number of part bindings differs between template and finalize.");
  }

  for (int i = 0; i < template_hard_binding.part_bindings_size(); ++i) {
    ABSL_RETURN_IF_ERROR(CheckIfPartBindingByteSizesMatched(
        template_hard_binding.part_bindings(i),
        finalize_hard_binding.part_bindings(i), proto_to_cbor_options));
  }
  return absl::OkStatus();
}

absl::StatusOr<bool> HasActiveClaimTimestamp(absl::string_view manifest_store) {
  ABSL_ASSIGN_OR_RETURN(ManifestStoreView view,
                        ManifestStoreView::Create(manifest_store));
  ABSL_ASSIGN_OR_RETURN(ManifestView manifest, view.GetActiveManifest());
  ABSL_ASSIGN_OR_RETURN(ClaimSignatureView claim_signature,
                        manifest.GetClaimSignature());
  ABSL_ASSIGN_OR_RETURN(CoseSign1TaggedStructure cose_sign1,
                        claim_signature.GetStruct());
  return cose_sign1.unprotected_header.sig_tst2.has_value();
}

}  // namespace

void ManifestStoreGeneratorImpl::GenerateManifestStore(
    GenerateManifestStoreRequest request, AsyncContext context,
    absl::AnyInvocable<void(absl::StatusOr<GenerateManifestStoreResponse>) &&>
        cb) const {
  int64_t max_manifest_store_size_bytes =
      request.options().max_manifest_store_size_bytes();
  if (max_manifest_store_size_bytes < 0) {
    AsyncResolver<GenerateManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(absl::InvalidArgumentError(
            "max_manifest_store_size_bytes must be non-negative"));
    return;
  }
  if (request.options().has_max_bytes_to_shift() &&
      request.options().max_bytes_to_shift() < 0) {
    AsyncResolver<GenerateManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(absl::InvalidArgumentError(
            "max_bytes_to_shift must be non-negative"));
    return;
  }
  if (max_manifest_store_size_bytes == 0) {
    max_manifest_store_size_bytes = kDefaultMaxManifestStoreSizeBytes;
  }
  auto request_ptr =
      std::make_shared<GenerateManifestStoreRequest>(std::move(request));
  auto manifest_or = manifest_assembler_->Assemble(*request_ptr);
  if (!manifest_or.ok()) {
    AsyncResolver<GenerateManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(manifest_or.status());
    return;
  }
  auto manifest_box = std::make_shared<ManifestBox>(std::move(*manifest_or));
  manifest_encoder_->Encode(
      manifest_box.get(), request_ptr->options(), context,
      [request_ptr, manifest_box, max_manifest_store_size_bytes,
       manifest_store_transformer = manifest_store_transformer_,
       cb = std::move(cb)](
          absl::StatusOr<ManifestEncoderResult> res_or) mutable {
        if (!res_or.ok()) {
          std::move(cb)(res_or.status());
          return;
        }
        absl::Cord encoded_store = res_or->box.Peek();
        absl::StatusOr<bool> is_timestamped =
            HasActiveClaimTimestamp(encoded_store.Flatten());
        if (!is_timestamped.ok()) {
          std::move(cb)(is_timestamped.status());
          return;
        }
        if (manifest_store_transformer != nullptr) {
          absl::Status transform_status =
              manifest_store_transformer->Transform(&res_or->box);
          if (!transform_status.ok()) {
            std::move(cb)(transform_status);
            return;
          }
        }
        absl::Cord manifest_store_cord = std::move(res_or->box).AsCord();
        if (manifest_store_cord.size() >
            static_cast<uint64_t>(max_manifest_store_size_bytes)) {
          std::move(cb)(absl::FailedPreconditionError(absl::Substitute(
              "Generated manifest store too large ($0 > $1)",
              manifest_store_cord.size(), max_manifest_store_size_bytes)));
          return;
        }
        GenerateManifestStoreResponse response;
        response.set_manifest_store(std::string(manifest_store_cord));
        response.set_is_timestamped(*is_timestamped);
        std::move(cb)(std::move(response));
      });
}

void ManifestStoreGeneratorImpl::GenerateManifestStoreTemplate(
    GenerateManifestStoreRequest request, AsyncContext context,
    absl::AnyInvocable<void(absl::StatusOr<ManifestStoreTemplate>) &&> cb)
    const {
  std::string generator_context = request.SerializeAsString();
  GenerateManifestStore(
      std::move(request), context,
      [generator_context = std::move(generator_context), cb = std::move(cb)](
          absl::StatusOr<GenerateManifestStoreResponse> generated_or) mutable {
        if (!generated_or.ok()) {
          std::move(cb)(generated_or.status());
          return;
        }
        ManifestStoreTemplate response;
        response.set_generator_context(std::move(generator_context));
        response.set_manifest_store_size(generated_or->manifest_store().size());
        std::move(cb)(std::move(response));
      });
}

void ManifestStoreGeneratorImpl::FinalizeManifestStore(
    FinalizeManifestStoreRequest request, AsyncContext context,
    absl::AnyInvocable<void(absl::StatusOr<FinalizeManifestStoreResponse>) &&>
        cb) const {
  if (request.manifest_store_template().generator_context().empty()) {
    AsyncResolver<FinalizeManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(absl::InvalidArgumentError(
            "The generator_context in the manifest store template is empty"));
    return;
  }
  if (request.manifest_store_template().manifest_store_size() < 0) {
    AsyncResolver<FinalizeManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(absl::InvalidArgumentError(
            "manifest_store_size in template must be non-negative"));
    return;
  }
  GenerateManifestStoreRequest generate_request;
  if (!generate_request.ParseFromString(
          request.manifest_store_template().generator_context())) {
    AsyncResolver<FinalizeManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(absl::InvalidArgumentError(
            "Failed to parse the generator_context in the manifest store "
            "template"));
    return;
  }
  if (generate_request.hard_binding().binding_case() !=
      request.hard_binding().binding_case()) {
    AsyncResolver<FinalizeManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(absl::InvalidArgumentError(
            "The hard binding in the finalize request must match the original "
            "request"));
    return;
  }
  absl::Status check_status = CheckIfHardBindingByteSizesMatched(
      generate_request.hard_binding(), request.hard_binding(),
      proto_to_cbor_options_);
  if (!check_status.ok()) {
    AsyncResolver<FinalizeManifestStoreResponse>(std::move(context),
                                                 std::move(cb))
        .Return(check_status);
    return;
  }
  generate_request.mutable_hard_binding()->Swap(request.mutable_hard_binding());

  size_t template_manifest_store_size =
      request.manifest_store_template().manifest_store_size();
  GenerateManifestStore(
      std::move(generate_request), context,
      [template_manifest_store_size, cb = std::move(cb)](
          absl::StatusOr<GenerateManifestStoreResponse> generated_or) mutable {
        if (!generated_or.ok()) {
          std::move(cb)(generated_or.status());
          return;
        }
        if (generated_or->manifest_store().size() !=
            template_manifest_store_size) {
          std::move(cb)(absl::InternalError(absl::Substitute(
              "The generated manifest store size ($0) does not "
              "match the estimated size ($1), even though the hard binding "
              "sizes match. This is an internal error.",
              generated_or->manifest_store().size(),
              template_manifest_store_size)));
          return;
        }
        FinalizeManifestStoreResponse response;
        response.mutable_manifest_store()->swap(
            *generated_or->mutable_manifest_store());
        response.set_is_timestamped(generated_or->is_timestamped());
        std::move(cb)(std::move(response));
      });
}

}  // namespace credentio