// 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.
//

// RIFF format information:
//   https://www.loc.gov/preservation/digital/formats/fdd/fdd000025.shtml
//   https://johnloomis.org/cpe102/asgn/asgn1/riff.html
//   https://www.tactilemedia.com/info/MCI_Control_Info.html
// C2PA embedding:
// https://spec.c2pa.org/specifications/specifications/2.4/specs/C2PA_Specification.html#_embedding_manifests_into_riff_based_assets
#include "formats/riff/extractor.h"

#include <cstdint>
#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/string_view.h"
#include "constants/labels.h"
#include "formats/asset_box.h"
#include "formats/byte_range.h"
#include "formats/riff/constants.h"
#include "formats/riff/reader.h"
#include "jumbf/utils.h"
#include "riegeli/bytes/reader.h"

namespace credentio {

namespace {

absl::StatusOr<std::vector<AssetBox>> ExtractBoxesInternal(
    riegeli::Reader& reader, int64_t end_offset) {
  std::vector<AssetBox> result;
  bool first = true;
  std::vector<uint64_t> container_ends;

  auto status = IterateOverRiffChunks(
      reader,
      [&result, &first, &container_ends,
       &reader](const RiffChunk& chunk) -> absl::StatusOr<bool> {
        while (!container_ends.empty() &&
               chunk.offset >= container_ends.back()) {
          container_ends.pop_back();
        }
        int depth = container_ends.size();

        if (depth == 0) {
          if (first) {
            if (chunk.id != kRiffChunkIdRiff) {
              return absl::InvalidArgumentError("RIFF chunk not found");
            }
            first = false;
          }

          result.push_back({
              .identifier = chunk.id,
              .byte_range = {.offset = chunk.offset, .length = 12},
          });

          if (chunk.HasSubchunks()) {
            container_ends.push_back(chunk.offset + chunk.length);
          }
          return true;
        }

        if (depth == 1) {
          if (chunk.id == kRiffChunkIdC2pa && IsRiffChunkTooLarge(chunk)) {
            return absl::InvalidArgumentError(
                "RIFF C2PA chunk is too large to extract");
          }

          result.push_back({
              .identifier = chunk.id,
              .byte_range = {.offset = chunk.offset, .length = chunk.length},
          });

          if (chunk.HasSubchunks()) {
            reader.Seek(chunk.offset + chunk.length);
          }
          return true;
        }

        return absl::InternalError("Unexpected depth in RIFF parser");
      },
      end_offset);

  ABSL_RETURN_IF_ERROR(status);
  return result;
}

// Returns the C2PA subchunk within the RIFF chunk, if any.
absl::StatusOr<std::optional<RiffChunk>> FindC2paChunk(riegeli::Reader& reader,
                                                       int64_t end_offset) {
  std::optional<RiffChunk> c2pa_chunk;
  RiffChunk riff_chunk;
  bool first = true;

  auto status = IterateOverRiffChunks(
      reader,
      [&c2pa_chunk, &first,
       &riff_chunk](const RiffChunk& chunk) -> absl::StatusOr<bool> {
        if (first) {
          if (chunk.id != kRiffChunkIdRiff) {
            return absl::InvalidArgumentError(
                "Input does not start with RIFF chunk");
          }
          riff_chunk = chunk;
          first = false;
          return true;
        }
        if (chunk.offset >= riff_chunk.offset + riff_chunk.length) {
          return false;  // Stop iterating, we finished the first RIFF chunk
        }
        if (chunk.id == kRiffChunkIdC2pa) {
          if (IsRiffChunkTooLarge(chunk)) {
            return absl::InvalidArgumentError(
                "RIFF C2PA chunk is too large to extract");
          }
          c2pa_chunk = chunk;
          return false;  // Found it, stop iterating
        }
        return true;
      },
      end_offset);

  ABSL_RETURN_IF_ERROR(status);
  if (first) {
    return absl::NotFoundError("No manifest store found");
  }
  return c2pa_chunk;
}

}  // namespace

absl::StatusOr<std::string> RiffExtractor::ExtractManifestStore(
    riegeli::Reader& input) const {
  ABSL_ASSIGN_OR_RETURN(std::optional<RiffChunk> c2pa_chunk,
                        FindC2paChunk(input, /*end_offset=*/-1));
  if (!c2pa_chunk.has_value()) {
    return absl::NotFoundError("No manifest store found");
  }
  int64_t offset = c2pa_chunk->data_offset;
  int64_t length = c2pa_chunk->data_length;
  if (!input.Seek(offset) || input.pos() != offset) {
    return absl::InvalidArgumentError(
        "The input could not be reset to the manifest store location.");
  }
  std::string result;
  if (!input.Read(length, result)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to read manifest store"));
  }
  return std::move(result);
}

absl::StatusOr<std::optional<ByteRange>>
RiffExtractor::ExtractManifestStoreLocation(riegeli::Reader& input,
                                            ExtractOptions options) const {
  ABSL_ASSIGN_OR_RETURN(std::optional<RiffChunk> c2pa_chunk,
                        FindC2paChunk(input, options.end_offset));
  if (!c2pa_chunk.has_value()) {
    if (options.requires_c2pa) {
      return absl::NotFoundError("No manifest store found");
    }
    return std::nullopt;
  }
  return ByteRange{
      .offset = c2pa_chunk->offset,
      .length = c2pa_chunk->length,
  };
}

absl::StatusOr<std::vector<AssetBox>> RiffExtractor::ExtractBoxes(
    riegeli::Reader& input, ExtractOptions options) const {
  ABSL_ASSIGN_OR_RETURN(std::vector<AssetBox> boxes,
                        ExtractBoxesInternal(input, options.end_offset));
  if (!options.requires_c2pa) {
    return std::move(boxes);
  }

  int64_t c2pa_count = 0;
  for (const auto& box : boxes) {
    if (box.identifier == kRiffChunkIdC2pa) {
      c2pa_count++;
      if (c2pa_count > 1) {
        return absl::NotFoundError("Multiple manifest stores found");
      }
    }
  }
  if (c2pa_count == 0) {
    return absl::NotFoundError("No manifest store found");
  }
  return std::move(boxes);
}

bool RiffExtractor::MightBeC2paManifestStore(absl::string_view payload) const {
  return jumbf::HasDescriptionBoxMatching(payload, kManifestStoreUuid,
                                          kMinimumJumbfDescriptionToggles,
                                          kManifestStoreLabel)
      .value_or(false);
}

}  // namespace credentio