// 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/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/jpeg/c2pa_segment_handler.h"
#include "formats/jpeg/reader.h"
#include "jumbf/utils.h"
#include "riegeli/bytes/reader.h"
#include "utils/byte_readers.h"

namespace credentio {
namespace {

// APP11 Marker and APP11 JPEG Extension Type values are sourced from
// Section A.4 JPEG XT boxes of the ISO/IEC 18477-3:2023
const uint16_t kApp11ExtensionType = 0x4a50;

}  // namespace

absl::StatusOr<std::string> JpegExtractor::ExtractManifestStore(
    riegeli::Reader& input) const {
  C2paSegmentHandler segment_handler;

  ABSL_RETURN_IF_ERROR(
      IterateOverJpegSegments(input, [&segment_handler](JpegSegment segment) {
        segment_handler.RecordSegment(std::move(segment));
        return true;
      }));

  ABSL_ASSIGN_OR_RETURN(auto c2pa_segment,
                        segment_handler.GetC2paSegment(input));

  return std::move(c2pa_segment.manifest_store);
}

absl::StatusOr<std::optional<ByteRange>>
JpegExtractor::ExtractManifestStoreLocation(riegeli::Reader& input,
                                            ExtractOptions options) const {
  C2paSegmentHandler c2pa_handler;
  ABSL_RETURN_IF_ERROR(IterateOverJpegSegments(
      input,
      [&c2pa_handler, &options, &input](JpegSegment segment) {
        if (options.end_offset >= 0 &&
            (segment.offset + segment.length) > options.end_offset) {
          // This segment starts or ends beyond the window, ignore, stop reading
          // and move the input's position to the declared end.
          if (!input.Seek(options.end_offset)) {
            return absl::StatusOr<bool>(input.StatusOrAnnotate(
                absl::DataLossError("Failed to seek to end offset")));
          }
          return absl::StatusOr<bool>(false);  // stop iteration
        }

        c2pa_handler.RecordSegment(std::move(segment));
        return absl::StatusOr<bool>(true);
      },
      options.end_offset));

  absl::StatusOr<C2paSegmentHandler::Segment> c2pa_segment =
      c2pa_handler.GetC2paSegment(input);
  if (!c2pa_segment.ok()) {
    if (options.requires_c2pa) {
      return c2pa_segment.status();
    }
    return std::nullopt;
  }

  return ByteRange{
      .offset = c2pa_segment->segment.offset,
      .length = c2pa_segment->segment.length,
  };
}

absl::StatusOr<std::vector<AssetBox>> JpegExtractor::ExtractBoxes(
    riegeli::Reader& input, ExtractOptions options) const {
  C2paSegmentHandler c2pa_handler;
  std::vector<AssetBox> boxes;
  ABSL_RETURN_IF_ERROR(IterateOverJpegSegments(
      input,
      [&boxes, &c2pa_handler](JpegSegment segment) {
        boxes.push_back(AssetBox{
            .identifier = segment.marker_label,
            .byte_range = {.offset = segment.offset, .length = segment.length},
        });

        c2pa_handler.RecordSegment(std::move(segment));
        return true;
      },
      options.end_offset));

  absl::StatusOr<C2paSegmentHandler::Segment> c2pa_segment =
      c2pa_handler.GetC2paSegment(input);
  if (!c2pa_segment.ok()) {
    if (options.requires_c2pa) {
      return c2pa_segment.status();
    }
    return boxes;
  }

  uint64_t c2pa_start = c2pa_segment->segment.offset;
  uint64_t c2pa_end = c2pa_start + c2pa_segment->segment.length;

  std::vector<AssetBox> merged_boxes;
  merged_boxes.reserve(boxes.size());
  for (const auto& box : boxes) {
    uint64_t box_start = box.byte_range.offset;
    uint64_t box_end = box_start + box.byte_range.length;

    if (box_end <= c2pa_start || box_start >= c2pa_end) {
      // Box is before or after the C2PA segment, add it.
      merged_boxes.push_back(std::move(box));
      continue;
    }

    if (box_end == c2pa_end) {
      // Box is the last C2PA Segment, add the box
      merged_boxes.push_back(AssetBox{
          .identifier = "C2PA",
          .byte_range = {.offset = c2pa_start,
                         .length = c2pa_segment->segment.length},
      });
      continue;
    }

    // Box is inside the C2PA segment, skip it.
  }
  return merged_boxes;
}

bool JpegExtractor::MightBeC2paManifestStore(absl::string_view payload) const {
  absl::StatusOr<uint16_t> extension_type = ConsumeUint<uint16_t>(&payload);
  if (!extension_type.ok() || *extension_type != kApp11ExtensionType) {
    return false;
  }

  // Segment ID
  auto iter_status = SkipBytes<uint16_t>(&payload);
  if (!iter_status.ok()) {
    return false;
  }

  // Sequence Number
  iter_status = SkipBytes<uint32_t>(&payload);
  if (!iter_status.ok()) {
    return false;
  }

  return jumbf::HasDescriptionBoxMatching(payload, kManifestStoreUuid,
                                          kMinimumJumbfDescriptionToggles,
                                          kManifestStoreLabel)
      .value_or(false);
}

}  // namespace credentio