// 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 "formats/riff/hasher_box.h"

#include <cstddef>
#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_cat.h"
#include "absl/strings/substitute.h"
#include "bindings/input_hasher.h"
#include "crypto/algorithms.h"
#include "formats/asset_box.h"
#include "formats/riff/extractor.h"
#include "proto/boxes_hash_assertion.pb.h"
#include "riegeli/bytes/reader.h"

namespace credentio {
namespace {

// Creator of a single BoxHash entry for a set of boxes.
// This isn't RIFF-specific; could move to a common location.
class BoxHasher {
 public:
  static absl::StatusOr<BoxHasher> Create(HashAlgorithm alg,
                                          BoxesHashAssertion* assertion) {
    ABSL_ASSIGN_OR_RETURN(std::unique_ptr<InputHasher> hasher,
                          InputHasher::Create(alg));
    return BoxHasher(std::move(hasher), assertion->add_boxes());
  }

  absl::Status HashBox(riegeli::Reader& input, const AssetBox& box) {
    entry_->add_names(box.identifier);
    ABSL_RETURN_IF_ERROR(
        hasher_->Update(input, box.byte_range.offset, box.byte_range.length));
    return absl::OkStatus();
  }

  void Finish() { entry_->set_hash(hasher_->Digest()); }

 private:
  BoxHasher(std::unique_ptr<InputHasher> hasher, BoxHash* entry)
      : hasher_(std::move(hasher)), entry_(entry) {}
  std::unique_ptr<InputHasher> hasher_;
  BoxHash* entry_;
};

void AddC2paEntry(BoxesHashAssertion& assertion) {
  BoxHash* entry = assertion.add_boxes();
  entry->add_names("C2PA");
  entry->set_hash(std::string(1, '\0'));
}

}  // namespace

absl::StatusOr<BoxesHashAssertion> DigestRiffBoxHash(riegeli::Reader& input,
                                                     HashAlgorithm alg) {
  BoxesHashAssertion assertion;
  assertion.set_alg(absl::StrCat(alg));
  ABSL_ASSIGN_OR_RETURN(
      std::vector<AssetBox> boxes,
      RiffExtractor().ExtractBoxes(input, {.requires_c2pa = false}));
  // Produce hashes for the boxes prior to the (existing or expected) C2PA box.
  std::optional<BoxHasher> hasher;
  size_t b = 0;
  for (; b < boxes.size(); ++b) {
    const auto& box = boxes[b];
    if (b == 0) {
      if (box.identifier != "RIFF") {
        return absl::InvalidArgumentError(absl::Substitute(
            "First box has ID '$0', expected 'RIFF'", box.identifier));
      }
    } else {
      if (box.identifier == "RIFF") {
        // Second RIFF box -- insert placeholder C2PA box first.
        break;
      }
    }
    if (box.identifier == "C2PA") {
      ++b;
      break;
    }
    if (!hasher) {
      ABSL_ASSIGN_OR_RETURN(hasher, BoxHasher::Create(alg, &assertion));
    }
    ABSL_RETURN_IF_ERROR(hasher->HashBox(input, box));
  }
  if (hasher) {
    hasher->Finish();
    hasher = std::nullopt;
  }
  if (b == 0) {
    return absl::InvalidArgumentError("RIFF chunk not found");
  }

  AddC2paEntry(assertion);

  // C2PA chunk should be the last box in the first RIFF chunk, so the next
  // chunk, if any, should be another top-level RIFF chunk.
  if (b < boxes.size() && boxes[b].identifier != "RIFF") {
    return absl::InvalidArgumentError(absl::Substitute(
        "Expected RIFF chunk, found '$0'", boxes[b].identifier));
  }

  // Produce hashes for any remaining boxes (after the C2PA chunk).
  for (; b < boxes.size(); ++b) {
    const auto& box = boxes[b];
    if (box.identifier == "C2PA") {
      return absl::InvalidArgumentError("Unexpected C2PA chunk found");
    }
    if (!hasher) {
      ABSL_ASSIGN_OR_RETURN(hasher, BoxHasher::Create(alg, &assertion));
    }
    ABSL_RETURN_IF_ERROR(hasher->HashBox(input, box));
  }
  if (hasher) {
    hasher->Finish();
    hasher = std::nullopt;
  }
  return assertion;
}

}  // namespace credentio