// 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.

#ifndef LIB_JXL_IMAGE_BUNDLE_H_
#define LIB_JXL_IMAGE_BUNDLE_H_

// The main image or frame consists of a bundle of associated images.

#include <jxl/cms_interface.h>
#include <jxl/memory_manager.h>

#include <cstddef>
#include <cstdint>
#include <memory>
#include <string>
#include <utility>
#include <vector>

#include "lib/jxl/base/common.h"
#include "lib/jxl/base/data_parallel.h"
#include "lib/jxl/base/rect.h"
#include "lib/jxl/base/status.h"
#include "lib/jxl/color_encoding_internal.h"
#include "lib/jxl/common.h"  // JPEGXL_ENABLE_TRANSCODE_JPEG
#include "lib/jxl/frame_header.h"
#include "lib/jxl/image.h"
#include "lib/jxl/image_metadata.h"
#include "lib/jxl/image_ops.h"
#include "lib/jxl/jpeg/jpeg_data.h"

namespace jxl {

// A bundle of color/alpha/depth/plane images.
class ImageBundle {
 public:
  // Uninitialized state for use as output parameter.
  explicit ImageBundle(JxlMemoryManager* memory_manager)
      : memory_manager_(memory_manager), metadata_(nullptr) {}
  // Caller is responsible for setting metadata before calling Set*.
  ImageBundle(JxlMemoryManager* memory_manager, const ImageMetadata* metadata)
      : memory_manager_(memory_manager), metadata_(metadata) {}

  // Move-only (allows storing in std::vector).
  ImageBundle(ImageBundle&&) = default;
  ImageBundle& operator=(ImageBundle&&) = default;

  StatusOr<ImageBundle> Copy() const {
    JxlMemoryManager* memory_manager = this->memory_manager();
    ImageBundle copy(memory_manager, metadata_);
    JXL_ASSIGN_OR_RETURN(
        copy.color_,
        Image3F::Create(memory_manager, color_.xsize(), color_.ysize()));
    JXL_RETURN_IF_ERROR(CopyImageTo(color_, &copy.color_));
    copy.c_current_ = c_current_;
    copy.extra_channels_.reserve(extra_channels_.size());
    for (const ImageF& plane : extra_channels_) {
      JXL_ASSIGN_OR_RETURN(
          ImageF ec,
          ImageF::Create(memory_manager, plane.xsize(), plane.ysize()));
      JXL_RETURN_IF_ERROR(CopyImageTo(plane, &ec));
      copy.extra_channels_.emplace_back(std::move(ec));
    }

    copy.jpeg_data =
        jpeg_data ? make_unique<jpeg::JPEGData>(*jpeg_data) : nullptr;
    copy.color_transform = color_transform;
    copy.chroma_subsampling = chroma_subsampling;

    return copy;
  }

  // -- SIZE

  size_t xsize() const {
    if (IsJPEG()) return jpeg_data->width;
    if (color_.xsize() != 0) return color_.xsize();
    return extra_channels_.empty() ? 0 : extra_channels_[0].xsize();
  }
  size_t ysize() const {
    if (IsJPEG()) return jpeg_data->height;
    if (color_.ysize() != 0) return color_.ysize();
    return extra_channels_.empty() ? 0 : extra_channels_[0].ysize();
  }
  Status ShrinkTo(size_t xsize, size_t ysize);

  // sizes taking orientation into account
  size_t oriented_xsize() const {
    if (static_cast<uint32_t>(metadata_->GetOrientation()) > 4) {
      return ysize();
    } else {
      return xsize();
    }
  }
  size_t oriented_ysize() const {
    if (static_cast<uint32_t>(metadata_->GetOrientation()) > 4) {
      return xsize();
    } else {
      return ysize();
    }
  }

  JxlMemoryManager* memory_manager_;

  // -- COLOR

  JxlMemoryManager* memory_manager() const { return memory_manager_; }

  // Whether color() is valid/usable. Returns true in most cases. Even images
  // with spot colors (one example of when !planes().empty()) typically have a
  // part that can be converted to RGB.
  bool HasColor() const { return color_.xsize() != 0; }

  // For resetting the size when switching from a reference to main frame.
  void RemoveColor() { color_ = Image3F(); }

  // Do not use if !HasColor().
  const Image3F& color() const {
    // If this fails, Set* was not called - perhaps because decoding failed?
    JXL_DASSERT(HasColor());
    return color_;
  }

  // Do not use if !HasColor().
  Image3F* color() {
    JXL_DASSERT(HasColor());
    return &color_;
  }

  // If c_current.IsGray(), all planes must be identical. NOTE: c_current is
  // independent of metadata()->color_encoding, which is the original, whereas
  // a decoder might return pixels in a different c_current.
  // This only sets the color channels, you must also make extra channels
  // match the amount that is in the metadata.
  Status SetFromImage(Image3F&& color, const ColorEncoding& c_current);

  // -- COLOR ENCODING

  const ColorEncoding& c_current() const { return c_current_; }

  // Returns whether the color image has identical planes. Once established by
  // Set*, remains unchanged until a subsequent Set* or TransformTo.
  bool IsGray() const { return c_current_.IsGray(); }

  bool IsSRGB() const { return c_current_.IsSRGB(); }
  bool IsLinearSRGB() const { return c_current_.IsLinearSRGB(); }

  // Set the c_current profile without doing any transformation, e.g. if the
  // transformation was already applied.
  void OverrideProfile(const ColorEncoding& new_c_current) {
    c_current_ = new_c_current;
  }

  // TODO(lode): TransformTo and CopyTo are implemented in enc_image_bundle.cc,
  // move these functions out of this header file and class, to
  // enc_image_bundle.h.

  // Transforms color to c_desired and sets c_current to c_desired. Alpha and
  // metadata remains unchanged.
  Status TransformTo(const ColorEncoding& c_desired, const JxlCmsInterface& cms,
                     ThreadPool* pool = nullptr);
  // Copies this:rect, converts to c_desired, and allocates+fills out.
  Status CopyTo(const Rect& rect, const ColorEncoding& c_desired,
                const JxlCmsInterface& cms, Image3F* out,
                ThreadPool* pool = nullptr) const;

  // Detect 'real' bit depth, which can be lower than nominal bit depth
  // (this is common in PNG), returns 'real' bit depth
  size_t DetectRealBitdepth() const;

  // -- ALPHA

  Status SetAlpha(ImageF&& alpha);
  bool HasAlpha() const {
    return metadata_->Find(ExtraChannel::kAlpha) != nullptr;
  }
  bool AlphaIsPremultiplied() const {
    const ExtraChannelInfo* eci = metadata_->Find(ExtraChannel::kAlpha);
    return (eci == nullptr) ? false : eci->alpha_associated;
  }
  const ImageF* alpha() const;
  ImageF* alpha();

  // -- EXTRA CHANNELS
  bool HasBlack() const {
    return metadata_->Find(ExtraChannel::kBlack) != nullptr;
  }
  const ImageF* black() const;

  // Extra channels of unknown interpretation (e.g. spot colors).
  Status SetExtraChannels(std::vector<ImageF>&& extra_channels);
  void ClearExtraChannels() { extra_channels_.clear(); }
  bool HasExtraChannels() const { return !extra_channels_.empty(); }
  const std::vector<ImageF>& extra_channels() const { return extra_channels_; }
  std::vector<ImageF>& extra_channels() { return extra_channels_; }

  const ImageMetadata* metadata() const { return metadata_; }

  Status VerifyMetadata() const;

  void SetDecodedBytes(size_t decoded_bytes) { decoded_bytes_ = decoded_bytes; }
  size_t decoded_bytes() const { return decoded_bytes_; }

  // -- JPEG transcoding:

  // Returns true if image does or will represent quantized DCT-8 coefficients,
  // stored in 8x8 pixel regions.
  bool IsJPEG() const {
#if JPEGXL_ENABLE_TRANSCODE_JPEG
    return jpeg_data != nullptr;
#else   // JPEGXL_ENABLE_TRANSCODE_JPEG
    return false;
#endif  // JPEGXL_ENABLE_TRANSCODE_JPEG
  }

  std::unique_ptr<jpeg::JPEGData> jpeg_data;
  // these fields are used to signal the input JPEG color space
  // NOTE: JPEG doesn't actually provide a way to determine whether YCbCr was
  // applied or not.
  ColorTransform color_transform = ColorTransform::kNone;
  YCbCrChromaSubsampling chroma_subsampling;

  FrameOrigin origin{0, 0};

  // Animation-related information, corresponding to the timecode and duration
  // fields of the jxl::AnimationFrame of the jxl::FrameHeader.
  // TODO(lode): ImageBundle is used here to carry the information from
  // jxl::FrameHeader, consider instead passing a jxl::FrameHeader directly to
  // EncodeFrame or having a field of that type here.
  uint32_t duration = 0;
  uint32_t timecode = 0;

  // TODO(lode): these fields do not match the JXL frame header, it should be
  // possible to specify up to 4 (3 if nonzero duration) slots to save this
  // frame as reference (see save_as_reference).
  bool use_for_next_frame = false;
  bool blend = false;
  BlendMode blendmode = BlendMode::kBlend;

  std::string name;

 private:
  // Called after any Set* to ensure their sizes are compatible.
  Status VerifySizes() const;

  // Required for TransformTo so that an ImageBundle is self-sufficient. Always
  // points to the same thing, but cannot be const-pointer because that prevents
  // the compiler from generating a move ctor.
  const ImageMetadata* metadata_;

  // Initialized by Set*:
  Image3F color_;  // If empty, planes_ is not; all planes equal if IsGray().
  ColorEncoding c_current_;  // of color_

  // Initialized by SetPlanes; size = ImageMetadata.num_extra_channels
  std::vector<ImageF> extra_channels_;

  // How many bytes of the input were actually read.
  size_t decoded_bytes_ = 0;
};

}  // namespace jxl

#endif  // LIB_JXL_IMAGE_BUNDLE_H_