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

#include <cstdint>
#include <ostream>
#include <string>

#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/str_format.h"
#include "riegeli/bytes/reader.h"
#include "riegeli/endian/endian_reading.h"

namespace credentio {
namespace {

const uint32_t kJpegStartMarker = 0xffd8;
const uint32_t kJpegEndMarker = 0xffd9;

// Converts a JPEG marker to a human readable label.
// Sourced from Table B.1 of ISO/IEC 10918-1:1994
std::string ConvertMarkerToLabel(uint16_t marker) {
  if (0xff00 != (marker & 0xff00) || marker == 0xffff) {
    // Markers are always 0xFFXX
    // 0xFFFF is not a valid JPEG marker
    return absl::StrFormat("%04x", marker);
  }
  if (0x00c0 == (marker & 0x00f0)) {
    switch (marker) {
      case 0xffc4:
        return "DHT";  // Define Huffman Table(s)
      case 0xffc8:
        return "JPG";  // Reserved for JPG extensions
      case 0xffcc:
        return "DAC";  // Define arithmetic coding conditioning(s)
      default:
        // Start Of Frame markers labeled numerically
        return absl::StrFormat("SOF%d", (marker & 0x000f));
    }
  }
  if (0x00d0 == (marker & 0x00f0)) {
    switch (marker) {
      case 0xffd8:
        return "SOI";  // Start of Image
      case 0xffd9:
        return "EOI";  // End of Image
      case 0xffda:
        return "SOS";  // Start of Scan
      case 0xffdb:
        return "DQT";  // Define Quantization Table(s)
      case 0xffdc:
        return "DNL";  // Define Number of Lines
      case 0xffdd:
        return "DRI";  // Define Restart Interval
      case 0xffde:
        return "DHP";  // Define Hierarchical Progression
      case 0xffdf:
        return "EXP";  // Expand reference component(s)
      default:
        // Restart with modulo 8 markers labeled numerically
        return absl::StrFormat("RST%d", (marker & 0x000f));
    }
  }
  if (0x00e0 == (marker & 0x00f0)) {
    // Reserved for APP markers labeled numerically
    return absl::StrFormat("APP%d", (marker & 0x000f));
  }
  if (0x00f0 == (marker & 0x00f0)) {
    switch (marker) {
      case 0xfffe:
        return "COM";  // Comment
      default:
        // Reserved for JPEG extensions labeled numerically
        return absl::StrFormat("JPG%d", (marker & 0x000f));
    }
  }
  if (marker == 0xff01) {
    return "TEM";  // For temporary private use in arithmetic coding
  }
  return "RES";  // Reserved
}
}  // namespace

std::string JpegBoxHeader::label() const {
  return ConvertMarkerToLabel(this->type);
}

absl::StatusOr<JpegBoxHeader> ConsumeJpegBoxHeader(riegeli::Reader& input) {
  JpegBoxHeader b = {.offset = input.pos()};
  if (!riegeli::ReadBigEndian<uint16_t>(input, b.type)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read type"));
  }

  if (0xff00 != (b.type & 0xff00) || b.type == 0xffff) {
    // Invalid marker.
    return absl::InvalidArgumentError(
        absl::StrFormat("Invalid JPEG marker: %d", b.type));
  }

  if (b.type == kJpegStartMarker || b.type == kJpegEndMarker ||
      b.type == 0xff01 || (b.type >= 0xffd0 && b.type <= 0xffd7)) {
    // No payload.
    b.size = sizeof(uint16_t);
    return b;
  }

  uint16_t chunk_size;
  if (!riegeli::ReadBigEndian<uint16_t>(input, chunk_size)) {
    return input.StatusOrAnnotate(absl::DataLossError("Failed to read size"));
  }
  if (chunk_size < 2) {
    return absl::InvalidArgumentError(
        absl::StrFormat("Invalid JPEG segment size: %d", chunk_size));
  }
  b.size = chunk_size + sizeof(b.type);
  return b;
}

void PrintTo(const JpegBoxHeader& x, ::std::ostream* os) {
  *os << absl::StrCat("{type: ", x.type, ", offset: ", x.offset,
                      ", size: ", x.size, ", label: \"", x.label(), "\"}");
}

}  // namespace credentio