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

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

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/substitute.h"
#include "formats/riff/chunk_header.h"
#include "riegeli/bytes/reader.h"

namespace credentio {

namespace {
constexpr uint64_t kMaxPayloadSize = 1024 * 1024 * 10;  // 10 MiB

absl::StatusOr<bool> IterateOverRiffChunksInternal(
    riegeli::Reader& reader, uint64_t end_offset,
    RiffChunkProcessor processor) {
  while (reader.pos() + ChunkHeader::header_size() <= end_offset) {
    uint64_t chunk_offset = reader.pos();
    ChunkHeader chunk_header;
    ABSL_RETURN_IF_ERROR(chunk_header.Read(reader));

    RiffChunk chunk = {
        .offset = chunk_offset,
        .length = chunk_header.chunk_size(),
        .id = chunk_header.id,
        .type = "",
        .data_offset = chunk_offset + chunk_header.header_size(),
        .data_length = chunk_header.data_size,
    };

    if (chunk.HasSubchunks()) {
      if (chunk.data_length < 4) {
        return absl::InvalidArgumentError(
            "RIFF/LIST chunk too short to contain a form type");
      }
      if (!reader.Read(4, chunk.type)) {
        return reader.StatusOrAnnotate(
            absl::DataLossError("kUnexpectedEof; type"));
      }
      chunk.data_offset += 4;
      chunk.data_length -= 4;
    }

    if (chunk.length < ChunkHeader::header_size()) {
      return absl::InvalidArgumentError("Invalid RIFF chunk size");
    }
    if (chunk.length > end_offset - chunk.offset) {
      return absl::InvalidArgumentError(
          "RIFF chunk extends beyond the end of the file");
    }
    uint64_t chunk_end_offset = chunk.offset + chunk.length;

    bool continue_processing;
    if (auto result = processor(chunk); result.ok()) {
      continue_processing = *result;
    } else {
      return result.status();
    }
    if (!continue_processing) {
      return false;
    }

    if (reader.pos() < chunk.data_offset + chunk.data_length &&
        chunk.HasSubchunks()) {
      if (auto sub_continue_or = IterateOverRiffChunksInternal(
              reader, chunk_end_offset, processor);
          sub_continue_or.ok()) {
        continue_processing = *sub_continue_or;
      } else {
        return sub_continue_or.status();
      }
      if (!continue_processing) {
        return false;
      }
    }

    reader.Seek(chunk_end_offset);
    if (reader.pos() <= chunk_offset) {
      return absl::InternalError(
          "RIFF parser detected a read stall (infinite loop prevented)");
    }
  }
  return true;
}

}  // namespace

absl::Status IterateOverRiffChunks(riegeli::Reader& reader,
                                   RiffChunkProcessor processor,
                                   int64_t end_offset) {
  if (!reader.SupportsSize() || !reader.Size().has_value()) {
    return absl::InvalidArgumentError(
        "manifest store not embedded: reader size cannot be determined");
  }
  uint64_t actual_end_offset =
      end_offset < 0 ? *reader.Size()
                     : std::min(static_cast<uint64_t>(*reader.Size()),
                                static_cast<uint64_t>(end_offset));

  if ((actual_end_offset - reader.pos()) == 0) {
    return absl::OkStatus();
  }

  return IterateOverRiffChunksInternal(reader, actual_end_offset, processor)
      .status();
}

absl::StatusOr<std::string> ReadRiffChunkData(riegeli::Reader& reader,
                                              const RiffChunk& chunk) {
  if (IsRiffChunkTooLarge(chunk)) {
    return absl::InvalidArgumentError(
        absl::Substitute("RIFF C2PA chunk is too large to extract ($0 > $1)",
                         chunk.data_length, kMaxPayloadSize));
  }

  if (reader.pos() != chunk.data_offset) {
    if (!reader.Seek(chunk.data_offset) || reader.pos() != chunk.data_offset) {
      return absl::InvalidArgumentError("Failed to seek to chunk data offset");
    }
  }

  std::string data;
  if (!reader.Read(chunk.data_length, data)) {
    return reader.StatusOrAnnotate(absl::DataLossError("kUnexpectedEof; data"));
  }
  return std::move(data);
}

bool IsRiffChunkTooLarge(const RiffChunk& chunk) {
  return chunk.data_length > kMaxPayloadSize;
}

}  // namespace credentio