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

#include "lib/extras/enc/exr.h"

#include <memory>

// IWYU pragma: no_include "ImathConfig.h"

// IWYU pragma: no_include "ImfChromaticities.h"

// IWYU pragma: no_include "ImfHeader.h"

// IWYU pragma: no_include "OpenEXRConfig.h"

#include "lib/extras/enc/encode.h"

#if !JPEGXL_ENABLE_EXR

namespace jxl {

namespace extras {

std::unique_ptr<Encoder> GetEXREncoder() { return nullptr; }

}  // namespace extras

}  // namespace jxl

#else  // JPEGXL_ENABLE_EXR

#include <ImathVec.h>

#include <ImfIO.h>

#include <ImfNamespace.h>  // IWYU pragma: keep

#include <ImfRgba.h>

#include <ImfRgbaFile.h>

#include <ImfStandardAttributes.h>

#include <ImfThreading.h>

#include <jxl/codestream_header.h>

#include <jxl/color_encoding.h>

#include <jxl/types.h>

#include <algorithm>

#include <cstddef>

#include <cstdint>

#include <cstring>

#include <utility>

#include <vector>

#include "lib/extras/packed_image.h"

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

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

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

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

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

namespace jxl {

namespace extras {

namespace {

namespace OpenEXR = OPENEXR_IMF_NAMESPACE;

namespace Imath = IMATH_NAMESPACE;

// OpenEXR::Int64 is deprecated in favor of using uint64_t directly, but using

// uint64_t as recommended causes build failures with previous OpenEXR versions

// on macOS, where the definition for OpenEXR::Int64 was actually not equivalent

// to uint64_t. This alternative should work in all cases.

using ExrInt64 = decltype(std::declval<OpenEXR::IStream>().tellg());

class InMemoryOStream : public OpenEXR::OStream {

 public:

  // `bytes` must outlive the InMemoryOStream.

  explicit InMemoryOStream(std::vector<uint8_t>* const bytes)

      : OStream(/*fileName=*/""), bytes_(*bytes) {}

  void write(const char c[], const int n) override {

    if (bytes_.size() < pos_ + n) {

      bytes_.resize(pos_ + n);

    }

    std::copy_n(c, n, bytes_.begin() + pos_);

    pos_ += n;

  }

  ExrInt64 tellp() override { return pos_; }

  void seekp(const ExrInt64 pos) override {

    if (bytes_.size() + 1 < pos) {

      bytes_.resize(pos - 1);

    }

    pos_ = pos;

  }

 private:

  std::vector<uint8_t>& bytes_;

  size_t pos_ = 0;

};

// Loads a Big-Endian float

float LoadBEFloat(const uint8_t* p) {

  uint32_t u = LoadBE32(p);

  float result;

  memcpy(&result, &u, 4);

  return result;

}

// Loads a Little-Endian float

float LoadLEFloat(const uint8_t* p) {

  uint32_t u = LoadLE32(p);

  float result;

  memcpy(&result, &u, 4);

  return result;

}

Status EncodeImageEXR(const PackedImage& image, const JxlBasicInfo& info,

                      const JxlColorEncoding& c_enc, ThreadPool* pool,

                      std::vector<uint8_t>* bytes) {

  OpenEXR::setGlobalThreadCount(0);

  const size_t xsize = image.xsize;

  const size_t ysize = image.ysize;

  const bool has_alpha = info.alpha_bits > 0;

  const bool alpha_is_premultiplied = FROM_JXL_BOOL(info.alpha_premultiplied);

  if (info.num_color_channels != 3) {

    return JXL_FAILURE("OpenEXR encoding: expected 3 color channels, got %u",

                       static_cast<unsigned>(info.num_color_channels));

  }

  if (c_enc.color_space != JXL_COLOR_SPACE_RGB &&

      c_enc.color_space != JXL_COLOR_SPACE_GRAY) {

    return JXL_FAILURE(

        "OpenEXR encoding: expected RGB (%d) or grayscale (%d) colorspace, got "

        "%d",

        static_cast<int>(JXL_COLOR_SPACE_RGB),

        static_cast<int>(JXL_COLOR_SPACE_GRAY),

        static_cast<int>(c_enc.color_space));

  }

  if (c_enc.transfer_function != JXL_TRANSFER_FUNCTION_LINEAR) {

    return JXL_FAILURE(

        "OpenEXR encoding: expected linear transfer function (%d), got %d",

        static_cast<int>(JXL_TRANSFER_FUNCTION_LINEAR),

        static_cast<int>(c_enc.transfer_function));

  }

  const size_t num_channels = 3 + (has_alpha ? 1 : 0);

  const JxlPixelFormat format = image.format;

  if (format.data_type != JXL_TYPE_FLOAT) {

    return JXL_FAILURE("Unsupported pixel format for OpenEXR output");

  }

  const uint8_t* in = static_cast<const uint8_t*>(image.pixels());

  size_t in_stride = num_channels * 4 * xsize;

  OpenEXR::Header header(xsize, ysize);

  if (c_enc.color_space == JXL_COLOR_SPACE_RGB) {

    OpenEXR::Chromaticities chromaticities;

    chromaticities.red =

        Imath::V2f(c_enc.primaries_red_xy[0], c_enc.primaries_red_xy[1]);

    chromaticities.green =

        Imath::V2f(c_enc.primaries_green_xy[0], c_enc.primaries_green_xy[1]);

    chromaticities.blue =

        Imath::V2f(c_enc.primaries_blue_xy[0], c_enc.primaries_blue_xy[1]);

    chromaticities.white =

        Imath::V2f(c_enc.white_point_xy[0], c_enc.white_point_xy[1]);

    OpenEXR::addChromaticities(header, chromaticities);

  }

  OpenEXR::addWhiteLuminance(header, info.intensity_target);

  auto loadFloat =

      format.endianness == JXL_BIG_ENDIAN ? LoadBEFloat : LoadLEFloat;

  auto loadAlpha =

      has_alpha ? loadFloat : [](const uint8_t* p) -> float { return 1.0f; };

  // Ensure that the destructor of RgbaOutputFile has run before we look at the

  // size of `bytes`.

  {

    InMemoryOStream os(bytes);

    OpenEXR::RgbaOutputFile output(

        os, header,

        c_enc.color_space == JXL_COLOR_SPACE_GRAY

            ? (has_alpha ? OpenEXR::WRITE_YA : OpenEXR::WRITE_Y)

            : (has_alpha ? OpenEXR::WRITE_RGBA : OpenEXR::WRITE_RGB));

    // How many rows to write at once. Again, the OpenEXR documentation

    // recommends writing the whole image in one call.

    const int y_chunk_size = ysize;

    std::vector<OpenEXR::Rgba> output_rows(xsize * y_chunk_size);

    for (size_t start_y = 0; start_y < ysize; start_y += y_chunk_size) {

      // Inclusive.

      const size_t end_y = std::min(start_y + y_chunk_size - 1, ysize - 1);

      output.setFrameBuffer(output_rows.data() - start_y * xsize,

                            /*xStride=*/1, /*yStride=*/xsize);

      for (size_t y = start_y; y <= end_y; ++y) {

        const uint8_t* in_row = &in[(y - start_y) * in_stride];

        OpenEXR::Rgba* const JXL_RESTRICT row_data =

            &output_rows[(y - start_y) * xsize];

        for (size_t x = 0; x < xsize; ++x) {

          const uint8_t* in_pixel = &in_row[4 * num_channels * x];

          float r = loadFloat(&in_pixel[0]);

          float g = loadFloat(&in_pixel[4]);

          float b = loadFloat(&in_pixel[8]);

          const float alpha = loadAlpha(&in_pixel[12]);

          if (!alpha_is_premultiplied) {

            r *= alpha;

            g *= alpha;

            b *= alpha;

          }

          row_data[x] = OpenEXR::Rgba(r, g, b, alpha);

        }

      }

      output.writePixels(/*numScanLines=*/end_y - start_y + 1);

    }

  }

  return true;

}

class EXREncoder : public Encoder {

  std::vector<JxlPixelFormat> AcceptedFormats() const override {

    std::vector<JxlPixelFormat> formats;

    for (const uint32_t num_channels : {3, 4}) {

      for (const JxlDataType data_type : {JXL_TYPE_FLOAT}) {

        for (JxlEndianness endianness : {JXL_BIG_ENDIAN, JXL_LITTLE_ENDIAN}) {

          formats.push_back(JxlPixelFormat{/*num_channels=*/num_channels,

                                           /*data_type=*/data_type,

                                           /*endianness=*/endianness,

                                           /*align=*/0});

        }

      }

    }

    return formats;

  }

  Status Encode(const PackedPixelFile& ppf, EncodedImage* encoded_image,

                ThreadPool* pool) const override {

    JXL_RETURN_IF_ERROR(VerifyBasicInfo(ppf.info));

    encoded_image->icc.clear();

    encoded_image->bitstreams.clear();

    encoded_image->bitstreams.reserve(ppf.frames.size());

    for (const auto& frame : ppf.frames) {

      JXL_RETURN_IF_ERROR(VerifyPackedImage(frame.color, ppf.info));

      encoded_image->bitstreams.emplace_back();

      JXL_RETURN_IF_ERROR(EncodeImageEXR(frame.color, ppf.info,

                                         ppf.color_encoding, pool,

                                         &encoded_image->bitstreams.back()));

    }

    return true;

  }

};

}  // namespace

std::unique_ptr<Encoder> GetEXREncoder() {

  return jxl::make_unique<EXREncoder>();

}

}  // namespace extras

}  // namespace jxl

#endif  // JPEGXL_ENABLE_EXR