/*

 * Copyright 2022 The Android Open Source Project

 *

 * 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

 *

 *      http://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.

 */

#ifndef ULTRAHDR_JPEGDECODERHELPER_H

#define ULTRAHDR_JPEGDECODERHELPER_H

#include <stdio.h>  // For jpeglib.h.

// C++ build requires extern C for jpeg internals.

#ifdef __cplusplus

extern "C" {

#endif

#include <jerror.h>

#include <jpeglib.h>

#ifdef __cplusplus

}  // extern "C"

#endif

#include <cstdint>

#include <memory>

#include <vector>

#include "ultrahdr_api.h"

namespace ultrahdr {

/*!\brief List of supported operations */

typedef enum {

  PARSE_STREAM = (1 << 0),   /**< Parse jpeg header, APPn markers (Exif, Icc, Xmp, Iso) */

  DECODE_STREAM = (1 << 16), /**< Single channel images are decoded to Grayscale format and multi

                                channel images are decoded to RGB format */

  DECODE_TO_YCBCR_CS = (1 << 17), /**< Decode image to YCbCr Color Space  */

  DECODE_TO_RGB_CS = (1 << 18),   /**< Decode image to RGB Color Space  */

} decode_mode_t;

/*!\brief Encapsulates a converter from JPEG to raw image format. This class is not thread-safe */

class JpegDecoderHelper {

 public:

  JpegDecoderHelper() = default;

  ~JpegDecoderHelper() = default;

  /*!\brief This function decodes the bitstream that is passed to it to the desired format and

   * stores the results internally. The result is accessible via getter functions.

   *

   * \param[in]  image    pointer to compressed image

   * \param[in]  length   length of compressed image

   * \param[in]  mode     output decode format

   *

   * \return uhdr_error_info_t #UHDR_CODEC_OK if operation succeeds, uhdr_codec_err_t otherwise.

   */

  uhdr_error_info_t decompressImage(const void* image, size_t length,

                                    decode_mode_t mode = DECODE_TO_YCBCR_CS);

  /*!\brief This function parses the bitstream that is passed to it and makes image information

   * available to the client via getter() functions. It does not decompress the image. That is done

   * by decompressImage().

   *

   * \param[in]  image    pointer to compressed image

   * \param[in]  length   length of compressed image

   *

   * \return uhdr_error_info_t #UHDR_CODEC_OK if operation succeeds, uhdr_codec_err_t otherwise.

   */

  uhdr_error_info_t parseImage(const void* image, size_t length) {

    return decompressImage(image, length, PARSE_STREAM);

  }

  /*! Below public methods are only effective if a call to decompressImage() is made and it returned

   * true. */

  /*!\brief returns decompressed image descriptor */

  uhdr_raw_image_t getDecompressedImage();

  /*!\brief returns pointer to decompressed image

   * \deprecated This function is deprecated instead use getDecompressedImage().

   */

  void* getDecompressedImagePtr() { return mResultBuffer.data(); }

  /*!\brief returns size of decompressed image

   * \deprecated This function is deprecated instead use getDecompressedImage().

   */

  size_t getDecompressedImageSize() { return mResultBuffer.size(); }

  /*! Below public methods are only effective if a call to parseImage() or decompressImage() is made

   * and it returned true. */

  /*!\brief returns image width */

  unsigned int getDecompressedImageWidth() { return mPlaneWidth[0]; }

  /*!\brief returns image height */

  unsigned int getDecompressedImageHeight() { return mPlaneHeight[0]; }

  /*!\brief returns number of components in image */

  unsigned int getNumComponentsInImage() { return mNumComponents; }

  /*!\brief returns pointer to xmp block present in input image */

  void* getXMPPtr() { return mXMPBuffer.data(); }

  /*!\brief returns size of xmp block */

  size_t getXMPSize() { return mXMPBuffer.size(); }

  /*!\brief returns pointer to exif block present in input image */

  void* getEXIFPtr() { return mEXIFBuffer.data(); }

  /*!\brief returns size of exif block */

  size_t getEXIFSize() { return mEXIFBuffer.size(); }

  /*!\brief returns pointer to icc block present in input image */

  void* getICCPtr() { return mICCBuffer.data(); }

  /*!\brief returns size of icc block */

  size_t getICCSize() { return mICCBuffer.size(); }

  /*!\brief returns pointer to iso block present in input image */

  void* getIsoMetadataPtr() { return mIsoMetadataBuffer.data(); }

  /*!\brief returns size of iso block */

  size_t getIsoMetadataSize() { return mIsoMetadataBuffer.size(); }

  /*!\brief returns the offset of exif data payload with reference to 'image' address that is passed

   * via parseImage()/decompressImage() call. Note this does not include jpeg marker (0xffe1) and

   * the next 2 bytes indicating the size of the payload. If exif block is not present in the image

   * passed, then it returns -1. */

  long getEXIFPos() { return mExifPayLoadOffset; }

 private:

  // max number of components supported

  static constexpr int kMaxNumComponents = 3;

  uhdr_error_info_t decode(const void* image, size_t length, decode_mode_t mode);

  uhdr_error_info_t decode(jpeg_decompress_struct* cinfo, uint8_t* dest);

  uhdr_error_info_t decodeToCSYCbCr(jpeg_decompress_struct* cinfo, uint8_t* dest);

  uhdr_error_info_t decodeToCSRGB(jpeg_decompress_struct* cinfo, uint8_t* dest);

  // temporary storage

  std::unique_ptr<uint8_t[]> mPlanesMCURow[kMaxNumComponents];

  std::vector<JOCTET> mResultBuffer;       // buffer to store decoded data

  std::vector<JOCTET> mXMPBuffer;          // buffer to store xmp data

  std::vector<JOCTET> mEXIFBuffer;         // buffer to store exif data

  std::vector<JOCTET> mICCBuffer;          // buffer to store icc data

  std::vector<JOCTET> mIsoMetadataBuffer;  // buffer to store iso data

  // image attributes

  uhdr_img_fmt_t mOutFormat;

  unsigned int mNumComponents;

  unsigned int mPlaneWidth[kMaxNumComponents];

  unsigned int mPlaneHeight[kMaxNumComponents];

  unsigned int mPlaneHStride[kMaxNumComponents];

  unsigned int mPlaneVStride[kMaxNumComponents];

  long mExifPayLoadOffset;  // Position of EXIF package, default value is -1 which means no EXIF

                            // package appears.

};

} /* namespace ultrahdr  */

#endif  // ULTRAHDR_JPEGDECODERHELPER_H