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

#include <algorithm>
#include <cstdint>
#include <optional>
#include <ostream>

#include "absl/status/status.h"
#include "absl/status/status_macros.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/str_cat.h"
#include "formats/jpeg/box_header.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/endian/endian_reading.h"

namespace credentio {
namespace {

constexpr uint16_t kApp11ExtensionType = 0x4a50;

// Returns the start of the next chunk following the scan data that starts at
// `offset`.
absl::StatusOr<int64_t> FindEndOfScan(riegeli::Reader& input, int64_t offset) {
  if (!input.Seek(offset)) {
    return input.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }

  uint8_t first = 0;
  uint8_t second = 0;
  while (true) {
    if (first == 0xff && second != 0) {
      break;
    }
    first = second;
    if (!input.ReadByte(second)) {
      return input.StatusOrAnnotate(absl::DataLossError("Failed to read byte"));
    }
  }
  return input.pos() - 2;  // Back up to start of chunk header.
}

absl::Status PopulateJumbfData(riegeli::Reader& reader, JpegSegment& segment) {
  if (!segment.app11_info.has_value()) {
    // Will never happen as it's always called from PopulateApp11Info
    return absl::InvalidArgumentError(
        "APP11 info not populated, cannot populate JUMBf data");
  }

  uint64_t header_size = 8;  // lbox + tbox
  if (segment.app11_info->payload_length < header_size) {
    // Not enough bytes for JUMBf header.
    return absl::OkStatus();
  }

  segment.app11_info->jumbf_data = JpegJumbfData{};
  if (!riegeli::ReadBigEndian<uint32_t>(reader,
                                        segment.app11_info->jumbf_data->lbox)) {
    return reader.StatusOrAnnotate(absl::DataLossError("Failed to read lbox"));
  }
  if (!riegeli::ReadBigEndian<uint32_t>(reader,
                                        segment.app11_info->jumbf_data->tbox)) {
    return reader.StatusOrAnnotate(absl::DataLossError("Failed to read tbox"));
  }
  if (segment.app11_info->jumbf_data->lbox == 1) {
    header_size += sizeof(uint64_t);

    if (segment.app11_info->payload_length < header_size) {
      return absl::InvalidArgumentError(
          "JUMBf box payload too short to contain xlbox");
    }
    if (!riegeli::ReadBigEndian<uint64_t>(
            reader, segment.app11_info->jumbf_data->xlbox)) {
      return reader.StatusOrAnnotate(
          absl::DataLossError("Failed to read xlbox"));
    }
  }

  uint64_t jumbf_size = segment.app11_info->jumbf_data->lbox == 1
                            ? segment.app11_info->jumbf_data->xlbox
                            : segment.app11_info->jumbf_data->lbox;
  if (jumbf_size < header_size) {
    return absl::InvalidArgumentError("JUMBf box payload too short");
  }

  segment.app11_info->jumbf_data->payload_offset =
      segment.app11_info->payload_offset + header_size;
  segment.app11_info->jumbf_data->payload_length =
      segment.app11_info->payload_length - header_size;
  return absl::OkStatus();
}

absl::Status PopulateApp11Info(riegeli::Reader& reader, JpegSegment& segment) {
  // APP11 segment payload must be at least 8 bytes long
  if (segment.payload_length < 8) {
    // Return OK without populating app11_info, treating it as non-C2PA.
    return absl::OkStatus();
  }

  // Seek to the start of the APP11 header.
  if (!reader.Seek(segment.payload_offset)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("Failed to seek to offset"));
  }

  segment.app11_info = JpegApp11Segment{};
  if (!riegeli::ReadBigEndian<uint16_t>(reader,
                                        segment.app11_info->extension_type)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("Failed to read extension type"));
  }
  if (!riegeli::ReadBigEndian<uint16_t>(reader,
                                        segment.app11_info->segment_id)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("Failed to read segment ID"));
  }
  if (!riegeli::ReadBigEndian<uint32_t>(reader,
                                        segment.app11_info->sequence_number)) {
    return reader.StatusOrAnnotate(
        absl::DataLossError("Failed to read sequence number"));
  }

  segment.app11_info->payload_offset = segment.payload_offset + 8;
  segment.app11_info->payload_length = segment.payload_length - 8;

  if (segment.app11_info->extension_type == kApp11ExtensionType) {
    return PopulateJumbfData(reader, segment);
  }
  return absl::OkStatus();
}

}  // namespace

absl::Status IterateOverJpegSegments(riegeli::Reader& reader,
                                     JpegSegmentProcessor processor,
                                     int64_t end_offset) {
  uint64_t current_offset = reader.pos();
  uint64_t actual_end_offset =
      end_offset == -1
          ? reader.Size().value_or(0)
          : std::min(static_cast<uint64_t>(reader.Size().value_or(0)),
                     static_cast<uint64_t>(end_offset));

  while (current_offset < actual_end_offset) {
    if (!reader.Seek(current_offset)) {
      return reader.StatusOrAnnotate(
          absl::DataLossError("Failed to seek to offset"));
    }

    JpegBoxHeader box_header;
    if (auto result = ConsumeJpegBoxHeader(reader); result.ok()) {
      box_header = *result;
    } else {
      return result.status();
    }

    JpegSegment segment = {
        .offset = box_header.offset,
        .marker = box_header.type,
        .marker_label = box_header.label(),
        .length = box_header.size,
        .payload_offset = box_header.size < 4
                              ? box_header.offset + box_header.size
                              : box_header.offset + 4,
        .payload_length = box_header.size < 4 ? 0 : box_header.size - 4,
    };

    if (segment.marker_label == "SOS" ||
        absl::StartsWith(segment.marker_label, "RST")) {
      // Chunk is followed by indefinite-length image data, which is
      // included in the box content for hashing purposes.
      uint64_t next_chunk_start = 0;
      if (auto result = FindEndOfScan(reader, segment.offset + segment.length);
          result.ok()) {
        next_chunk_start = *result;
      } else {
        return result.status();
      }
      segment.length = next_chunk_start - segment.offset;
      segment.payload_length = next_chunk_start - segment.payload_offset;
      // Return to the start of the payload.
      if (!reader.Seek(segment.payload_offset)) {
        return reader.StatusOrAnnotate(
            absl::DataLossError("Failed to seek to offset"));
      }
    }

    if (segment.marker_label == "APP11") {
      ABSL_RETURN_IF_ERROR(PopulateApp11Info(reader, segment));
    }

    if (segment.length > actual_end_offset - segment.offset &&
        segment.payload_length > 0) {
      // This box is truncated.
      return absl::InvalidArgumentError("truncated JPEG box");
    }
    current_offset = segment.offset + segment.length;

    bool continue_processing = true;
    if (auto result = processor(segment); result.ok()) {
      continue_processing = *result;
    } else {
      return result.status();
    }
    if (!continue_processing) {
      return absl::OkStatus();
    }

    if (segment.marker_label == "EOI") {
      break;
    }
  }

  if (current_offset < actual_end_offset) {
    // End of first image, any additional data is just a `c2pa.after` segment.
    JpegSegment after_segment =
        JpegSegment{.offset = current_offset,
                    .marker = 0x0000,
                    .marker_label = "c2pa.after",
                    .length = actual_end_offset - current_offset,
                    .payload_offset = 0,
                    .payload_length = 0};
    return processor(after_segment).status();
  }

  return absl::OkStatus();
}

void PrintTo(const JpegJumbfData& x, ::std::ostream* os) {
  *os << absl::StrCat("{lbox: ", x.lbox, ", tbox: ", x.tbox,
                      ", xlbox: ", x.xlbox,
                      ", payload_offset: ", x.payload_offset,
                      ", payload_length: ", x.payload_length, "}");
}

void PrintTo(const JpegApp11Segment& x, ::std::ostream* os) {
  *os << "{extension_type: " << x.extension_type
      << ", segment_id: " << x.segment_id
      << ", sequence_number: " << x.sequence_number
      << ", payload_offset: " << x.payload_offset
      << ", payload_length: " << x.payload_length << ", jumbf_data: ";
  if (x.jumbf_data.has_value()) {
    PrintTo(*x.jumbf_data, os);
  } else {
    *os << "nullopt";
  }
  *os << "}";
}

void PrintTo(const JpegSegment& x, ::std::ostream* os) {
  *os << "{offset: " << x.offset << ", marker: " << x.marker
      << ", marker_label: \"" << x.marker_label << "\""
      << ", length: " << x.length << ", payload_offset: " << x.payload_offset
      << ", payload_length: " << x.payload_length << ", app11_info: ";
  if (x.app11_info.has_value()) {
    PrintTo(*x.app11_info, os);
  } else {
    *os << "nullopt";
  }
  *os << "}";
}

}  // namespace credentio