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

#include <cstddef>
#include <cstdint>

#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 "riegeli/bytes/reader.h"
#include "riegeli/bytes/string_reader.h"
#include "riegeli/bytes/writer.h"
#include "riegeli/endian/endian_reading.h"
#include "riegeli/endian/endian_writing.h"

namespace credentio {
namespace {

constexpr uint16_t kMaxPayloadSegmentByteSize = 65000;

struct JumbfHeader {
  uint32_t lbox;
  uint32_t tbox;
  uint64_t xlbox;

  uint64_t byte_size() const {
    return sizeof(lbox) + sizeof(tbox) + (lbox == 1 ? sizeof(xlbox) : 0);
  }
  uint64_t box_size() const { return lbox == 1 ? xlbox : lbox; }
};

absl::StatusOr<JumbfHeader> ConsumePayloadJumbfHeader(riegeli::Reader& reader) {
  JumbfHeader header;
  if (!riegeli::ReadBigEndian<uint32_t>(reader, header.lbox)) {
    return reader.StatusOrAnnotate(absl::DataLossError("Failed to read lbox"));
  }
  if (!riegeli::ReadBigEndian<uint32_t>(reader, header.tbox)) {
    return reader.StatusOrAnnotate(absl::DataLossError("Failed to read tbox"));
  }
  if (header.lbox == 1) {
    if (!riegeli::ReadBigEndian<uint64_t>(reader, header.xlbox)) {
      return reader.StatusOrAnnotate(
          absl::DataLossError("Failed to read xlbox"));
    }
  }
  return header;
}

absl::Status WritePayloadChunk(const JumbfHeader& header,
                               absl::string_view payload_chunk,
                               uint16_t segment_id, uint32_t sequence_num,
                               riegeli::Writer& output) {
  if (!riegeli::WriteBigEndian<uint16_t>(0xFFEB, output)) {
    return output.status();
  }
  uint64_t total_seg_length =
      sizeof(uint16_t) /*seg_length*/ + sizeof(uint16_t) /*ext_type*/ +
      sizeof(uint16_t) /*segment_id*/ + sizeof(uint32_t) /*sequence_num*/ +
      header.byte_size() + payload_chunk.size();
  if (total_seg_length > UINT16_MAX) {
    return absl::OutOfRangeError("maximum segment size too large");
  }
  uint16_t seg_length = static_cast<uint16_t>(total_seg_length);
  if (!riegeli::WriteBigEndian<uint16_t>(seg_length, output)) {
    return output.status();
  }
  // JPEG extension type.
  if (!riegeli::WriteBigEndian<uint16_t>(0x4A50, output)) {
    return output.status();
  }
  if (!riegeli::WriteBigEndian<uint16_t>(segment_id, output)) {
    return output.status();
  }
  if (!riegeli::WriteBigEndian<uint32_t>(sequence_num, output)) {
    return output.status();
  }
  if (!riegeli::WriteBigEndian<uint32_t>(header.lbox, output)) {
    return output.status();
  }
  if (!riegeli::WriteBigEndian<uint32_t>(header.tbox, output)) {
    return output.status();
  }
  if (header.lbox == 1) {
    if (!riegeli::WriteBigEndian<uint64_t>(header.xlbox, output)) {
      return output.status();
    }
  }
  if (!output.Write(payload_chunk)) {
    return output.status();
  }
  return absl::OkStatus();
}
}  // namespace

uint64_t CalculateTotalJpegApp11SegmentsSize(uint64_t jumbf_payload_size) {
  uint64_t jpeg_header_size = sizeof(uint16_t) +  // marker
                              sizeof(uint16_t) +  // size
                              sizeof(uint16_t) +  // ext_type
                              sizeof(uint16_t) +  // segment_id
                              sizeof(uint32_t);   // sequence_num
  uint64_t jumbf_header_size =
      JumbfHeader{.lbox = static_cast<uint32_t>(
                      jumbf_payload_size > UINT32_MAX ? 1 : 0)}
          .byte_size();

  uint64_t payload_size_without_header = jumbf_payload_size - jumbf_header_size;
  uint64_t number_of_chunks =
      (payload_size_without_header / kMaxPayloadSegmentByteSize) +
      (payload_size_without_header % kMaxPayloadSegmentByteSize != 0);

  return ((jpeg_header_size + jumbf_header_size) * number_of_chunks) +
         payload_size_without_header;
}

absl::Status WriteJpegApp11Payload(uint16_t segment_id,
                                   absl::string_view payload,
                                   riegeli::Writer& output) {
  riegeli::StringReader<> payload_header_reader(payload.substr(0, 16));
  ABSL_ASSIGN_OR_RETURN(auto header,
                        ConsumePayloadJumbfHeader(payload_header_reader));
  if (header.box_size() != payload.size()) {
    return absl::InvalidArgumentError(absl::Substitute(
        "invalid jumbf box -- box size does not match payload size: $0 vs $1",
        header.box_size(), payload.size()));
  }
  uint32_t seq_num = 1;  // Starts at 1 per ISO/IEC 19566-5.
  size_t cursor = header.byte_size();
  while (cursor < payload.size()) {
    uint16_t chunk_size = kMaxPayloadSegmentByteSize;
    ABSL_RETURN_IF_ERROR(WritePayloadChunk(header,
                                           payload.substr(cursor, chunk_size),
                                           segment_id, seq_num, output));
    cursor += chunk_size;
    if (++seq_num == 0) {
      return absl::OutOfRangeError(absl::Substitute(
          "too many segments required for payload of size $0", payload.size()));
    }
  }
  return absl::OkStatus();
}
}  // namespace credentio