// Copyright (c) the JPEG XL Project Authors. All rights reserved.

//

// Use of this source code is governed by a BSD-style

// license that can be found in the LICENSE file.

// Helper class for storing external (int or float, interleaved) images. This is

// the common format used by other libraries and in the libjxl API.

#include "packed_image.h"

#include <jxl/codestream_header.h>

#include <jxl/encode.h>

#include <jxl/types.h>

#include <algorithm>

#include <cstddef>

#include <cstdint>

#include <cstdlib>

#include <cstring>

#include <functional>

#include <memory>

#include <string>

#include <utility>

#include <vector>

#include "lib/jxl/base/byte_order.h"

#include "lib/jxl/base/common.h"

#include "lib/jxl/base/status.h"

namespace jxl {

namespace extras {

// Class representing an interleaved image with a bunch of channels.

StatusOr<PackedImage> PackedImage::Create(size_t xsize, size_t ysize,

                                          const JxlPixelFormat& format) {

  JXL_ASSIGN_OR_RETURN(size_t stride, CalcStride(format, xsize));

  size_t pixels_size = ysize * stride;

  if ((pixels_size / stride) != ysize) {

    return JXL_FAILURE("Image too big");

  }

  PackedImage image(xsize, ysize, format, stride);

  if (!image.pixels()) {

    // TODO(szabadka): use specialized OOM error code

    return JXL_FAILURE("Failed to allocate memory for image");

  }

  return image;

}

StatusOr<PackedImage> PackedImage::Copy() const {

  // Resulting copy_stride have to be less or equal to original -> always ok.

  JXL_ASSIGN_OR_RETURN(size_t copy_stride, CalcStride(format, xsize));

  PackedImage copy(xsize, ysize, format, copy_stride);

  const uint8_t* orig_pixels = reinterpret_cast<const uint8_t*>(pixels());

  uint8_t* copy_pixels = reinterpret_cast<uint8_t*>(copy.pixels());

  if (stride == copy_stride) {

    // Same stride -> copy in one go.

    memcpy(copy_pixels, orig_pixels, ysize * stride);

  } else {

    // Otherwise, copy row-wise.

    JXL_DASSERT(copy_stride < stride);

    for (size_t y = 0; y < ysize; ++y) {

      memcpy(copy_pixels + y * copy_stride, orig_pixels + y * stride,

             copy_stride);

    }

  }

  return copy;

}

Status PackedImage::ValidateDataType(JxlDataType data_type) {

  if ((data_type != JXL_TYPE_UINT8) && (data_type != JXL_TYPE_UINT16) &&

      (data_type != JXL_TYPE_FLOAT) && (data_type != JXL_TYPE_FLOAT16)) {

    return JXL_FAILURE("Unhandled data type: %d", static_cast<int>(data_type));

  }

  return true;

}

size_t PackedImage::BitsPerChannel(JxlDataType data_type) {

  switch (data_type) {

    case JXL_TYPE_UINT8:

      return 8;

    case JXL_TYPE_UINT16:

      return 16;

    case JXL_TYPE_FLOAT:

      return 32;

    case JXL_TYPE_FLOAT16:

      return 16;

    default:

      JXL_DEBUG_ABORT("Unreachable");

      return 0;

  }

}

// Logical resize; use Copy() for storage reallocation, if necessary.

Status PackedImage::ShrinkTo(size_t new_xsize, size_t new_ysize) {

  if (new_xsize > xsize || new_ysize > ysize) {

    return JXL_FAILURE("Cannot shrink PackedImage to a larger size");

  }

  xsize = new_xsize;

  ysize = new_ysize;

  return true;

}

PackedImage::PackedImage(size_t xsize, size_t ysize,

                         const JxlPixelFormat& format, size_t stride)

    : xsize(xsize),

      ysize(ysize),

      stride(stride),

      format(format),

      pixels_size(ysize * stride),

      pixels_(malloc(std::max<size_t>(1, pixels_size)), free) {

  bytes_per_channel_ = BitsPerChannel(format.data_type) / jxl::kBitsPerByte;

  pixel_stride_ = format.num_channels * bytes_per_channel_;

  swap_endianness_ = SwapEndianness(format.endianness);

}

StatusOr<size_t> PackedImage::CalcStride(const JxlPixelFormat& format,

                                         size_t xsize) {

  size_t multiplier = (BitsPerChannel(format.data_type) * format.num_channels /

                       jxl::kBitsPerByte);

  size_t stride;

  if (!SafeMul(xsize, multiplier, stride)) {

    return JXL_FAILURE("Image too big");

  }

  if (!SafeRoundUpTo(stride, format.align, stride)) {

    return JXL_FAILURE("Image too big");

  }

  return stride;

}

PackedFrame::PackedFrame(PackedImage&& image) : color(std::move(image)) {}

PackedFrame::PackedFrame(PackedFrame&& other) = default;

PackedFrame& PackedFrame::operator=(PackedFrame&& other) = default;

PackedFrame::~PackedFrame() = default;

StatusOr<PackedFrame> PackedFrame::Create(size_t xsize, size_t ysize,

                                          const JxlPixelFormat& format) {

  JXL_ASSIGN_OR_RETURN(PackedImage image,

                       PackedImage::Create(xsize, ysize, format));

  PackedFrame frame(std::move(image));

  return frame;

}

StatusOr<PackedFrame> PackedFrame::Copy() const {

  JXL_ASSIGN_OR_RETURN(

      PackedFrame copy,

      PackedFrame::Create(color.xsize, color.ysize, color.format));

  copy.frame_info = frame_info;

  copy.name = name;

  JXL_ASSIGN_OR_RETURN(copy.color, color.Copy());

  for (const auto& ec : extra_channels) {

    JXL_ASSIGN_OR_RETURN(PackedImage ec_copy, ec.Copy());

    copy.extra_channels.emplace_back(std::move(ec_copy));

  }

  return copy;

}

// Logical resize; use Copy() for storage reallocation, if necessary.

Status PackedFrame::ShrinkTo(size_t new_xsize, size_t new_ysize) {

  JXL_RETURN_IF_ERROR(color.ShrinkTo(new_xsize, new_ysize));

  for (auto& ec : extra_channels) {

    JXL_RETURN_IF_ERROR(ec.ShrinkTo(new_xsize, new_ysize));

  }

  frame_info.layer_info.xsize = new_xsize;

  frame_info.layer_info.ysize = new_ysize;

  return true;

}

ChunkedPackedFrame::ChunkedPackedFrame(

    size_t xsize, size_t ysize,

    std::function<JxlChunkedFrameInputSource()> get_input_source)

    : xsize(xsize),

      ysize(ysize),

      get_input_source_(std::move(get_input_source)) {

  const auto input_source = get_input_source_();

  input_source.get_color_channels_pixel_format(input_source.opaque, &format);

}

PackedPixelFile::PackedPixelFile() { JxlEncoderInitBasicInfo(&info); };

Status PackedPixelFile::ShrinkTo(size_t new_xsize, size_t new_ysize) {

  for (auto& frame : frames) {

    JXL_RETURN_IF_ERROR(frame.ShrinkTo(new_xsize, new_ysize));

  }

  info.xsize = new_xsize;

  info.ysize = new_ysize;

  return true;

}

}  // namespace extras

}  // namespace jxl