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

#include <cstddef>
#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/string_view.h"
#include "absl/strings/substitute.h"
#include "absl/types/span.h"
#include "constants/labels.h"
#include "formats/jpeg/reader.h"
#include "jumbf/box.h"
#include "jumbf/constants.h"
#include "jumbf/parse.h"
#include "riegeli/bytes/reader.h"

namespace credentio {
namespace {

constexpr uint64_t kMaxPayloadSize = 1024 * 1024 * 10;  // 10 MiB
constexpr size_t kMaxSegmentsPerId = 1000;

bool IsValidNextSegment(const JpegSegment& prev_segment,
                        const JpegSegment& curr_segment) {
  if (!prev_segment.app11_info.has_value() ||
      !curr_segment.app11_info.has_value()) {
    // Missing APP11 information, invalid.
    return false;
  }
  if (!prev_segment.app11_info->jumbf_data.has_value() ||
      !curr_segment.app11_info->jumbf_data.has_value()) {
    // Segments do not contain JUMBF data, invalid.
    return false;
  }

  // C2PA Spec requires sequence numbers to be increasing.
  if (prev_segment.app11_info->sequence_number >=
      curr_segment.app11_info->sequence_number) {
    // Sequence number is not increasing, invalid.
    return false;
  }

  // C2PA Spec requires segments to be contiguous.
  if ((prev_segment.offset + prev_segment.length) != curr_segment.offset) {
    // Segments are not contiguous, invalid.
    return false;
  }

  // All segments must have the same lbox and xlbox.
  if (prev_segment.app11_info->jumbf_data->lbox !=
      curr_segment.app11_info->jumbf_data->lbox) {
    // Lbox is different, invalid.
    return false;
  }
  if (prev_segment.app11_info->jumbf_data->xlbox !=
      curr_segment.app11_info->jumbf_data->xlbox) {
    // Xlbox is different, invalid.
    return false;
  }

  return true;
}

absl::StatusOr<std::string> ExtractCompletePayload(
    riegeli::Reader& input, absl::Span<const JpegSegment> segments) {
  if (segments.empty() || !segments.front().app11_info.has_value() ||
      !segments.front().app11_info->jumbf_data.has_value()) {
    return "";
  }
  uint64_t expected_payload_size =
      segments.front().app11_info->jumbf_data->lbox == 1
          ? segments.front().app11_info->jumbf_data->xlbox
          : segments.front().app11_info->jumbf_data->lbox;
  if (expected_payload_size == 0 || expected_payload_size > kMaxPayloadSize) {
    return absl::InvalidArgumentError(absl::Substitute(
        "Declared JUMBF payload size is invalid ($0)", expected_payload_size));
  }

  std::string payload;
  for (const JpegSegment& segment : segments) {
    if (!segment.app11_info.has_value() ||
        !segment.app11_info->jumbf_data.has_value()) {
      // Segment does not contain JUMBF data within an APP11 segment, skip it.
      continue;
    }

    uint64_t payload_offset = segment.app11_info->payload_offset;
    uint64_t payload_length = segment.app11_info->payload_length;
    if (!payload.empty()) {
      // This is not the first segment, so skip the JUMBF header data
      payload_offset = segment.app11_info->jumbf_data->payload_offset;
      payload_length = segment.app11_info->jumbf_data->payload_length;
    }

    if (payload.size() + payload_length > kMaxPayloadSize) {
      return absl::InvalidArgumentError(
          absl::Substitute("JPEG segment is too large to extract ($0 > $1)",
                           payload.size() + payload_length, kMaxPayloadSize));
    }

    if (!input.Seek(payload_offset) || input.pos() != payload_offset) {
      return absl::InvalidArgumentError("Failed to seek to payload offset");
    }

    std::string segment_payload;
    if (!input.Read(payload_length, segment_payload)) {
      return input.StatusOrAnnotate(
          absl::DataLossError("Failed to read payload"));
    }
    payload.append(segment_payload);
  }

  if (payload.size() != expected_payload_size) {
    // Incomplete payload, return empty string.
    return "";
  }

  return payload;
}

bool IsC2paManifestStore(absl::string_view payload) {
  absl::StatusOr<jumbf::SuperBox> box =
      jumbf::ConsumeSuperBox(&payload, /*recursion_limit=*/0);
  if (!box.ok()) {
    // If we can't read a JUMBF Superbox, it's not a C2PA Manifest Store,
    // ignore.
    return false;
  }

  if (box->description.type_uuid != kManifestStoreUuid) {
    // JUMBF Superbox is not a C2PA Manifest Store, ignore.
    return false;
  };

  if (!box->description.requestable) {
    // Manifest store must be requestable.
    return false;
  }

  return box->description.label == kManifestStoreLabel;
}

}  // namespace

void C2paSegmentHandler::RecordSegment(JpegSegment segment) {
  if (!segment.app11_info.has_value() ||
      !segment.app11_info->jumbf_data.has_value()) {
    // Segment does not contain APP11 information, skip it.
    return;
  }

  const JpegApp11Segment& app11_info = *segment.app11_info;
  const JpegJumbfData& jumbf_data = *app11_info.jumbf_data;

  if (ignored_segment_ids_.contains(app11_info.segment_id)) {
    // Already ignored this segment ID, skip it.
    return;
  }

  if (app11_info.segment_id == 0) {
    // Invalid segment id, skip it.
    return;
  }

  if (jumbf_data.tbox != jumbf::kSuperBoxType) {
    // Segment is not a JUMBF Superbox, skip it.
    return;
  }

  auto [it, inserted] = segment_map_.try_emplace(app11_info.segment_id);
  if (inserted) {
    // First segment for this segment ID, record it.
    it->second.push_back(std::move(segment));
  } else {
    // Segment ID already exists.
    if (it->second.size() >= kMaxSegmentsPerId ||
        !IsValidNextSegment(it->second.back(), segment)) {
      // Exceeded max segments per ID or invalid continuing segment, remove from
      // map and ignore.
      segment_map_.erase(it);
      ignored_segment_ids_.insert(app11_info.segment_id);
    } else {
      // Valid continuing segment, record it.
      it->second.push_back(std::move(segment));
    }
  }
}

absl::StatusOr<C2paSegmentHandler::Segment> C2paSegmentHandler::GetC2paSegment(
    riegeli::Reader& input) const {
  C2paSegmentHandler::Segment result;
  int64_t c2pa_count = 0;

  for (const auto& [segment_id, segments] : segment_map_) {
    if (segments.empty()) {
      continue;
    }
    ABSL_ASSIGN_OR_RETURN(std::string payload,
                          ExtractCompletePayload(input, segments));
    if (payload.empty()) {
      // Incomplete payload, ignore.
      continue;
    }

    if (IsC2paManifestStore(payload)) {
      ++c2pa_count;
      if (c2pa_count > 1) {
        return absl::NotFoundError("Multiple manifest stores found");
      }

      result.manifest_store = std::move(payload);
      result.segment = JpegSegment{
          .offset = segments.front().offset,
          .marker_label = "C2PA",
          .length = (segments.back().offset - segments.front().offset) +
                    segments.back().length,
      };
    }
  }

  if (c2pa_count == 0) {
    return absl::NotFoundError("No manifest store found");
  }
  return std::move(result);
}

}  // namespace credentio