// 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 <jxl/codestream_header.h>

#include <jxl/color_encoding.h>

#include <jxl/decode.h>

#include <jxl/decode_cxx.h>

#include <jxl/encode.h>

#include <jxl/encode_cxx.h>

#include <jxl/thread_parallel_runner.h>

#include <jxl/thread_parallel_runner_cxx.h>

#include <jxl/types.h>

#include <cstdint>

#include <cstdlib>

#include <cstring>

#include <vector>

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

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

#include "lib/jxl/fuzztest.h"

#include "tools/tracking_memory_manager.h"

namespace {

using ::jpegxl::tools::kGiB;

using ::jpegxl::tools::TrackingMemoryManager;

using ::jxl::Status;

using ::jxl::StatusOr;

void Check(bool ok) {

  if (!ok) {

    JXL_CRASH();

  }

}

struct FuzzSpec {

  uint32_t xsize;

  uint32_t ysize;

  bool grayscale;

  bool alpha;

  uint8_t bit_depth;  // 1 - 16

  struct IntOptionSpec {

    JxlEncoderFrameSettingId flag;

    int min;

    int max;

    int value;

  };

  std::vector<IntOptionSpec> int_options = {

      IntOptionSpec{JXL_ENC_FRAME_SETTING_EFFORT, 1, 9, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_DECODING_SPEED, 0, 4, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_NOISE, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_DOTS, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_PATCHES, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_EPF, -1, 3, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_GABORISH, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_MODULAR, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_KEEP_INVISIBLE, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_RESPONSIVE, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_PROGRESSIVE_AC, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_QPROGRESSIVE_AC, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_PROGRESSIVE_DC, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_PALETTE_COLORS, -1, 255, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_LOSSY_PALETTE, -1, 1, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_COLOR_TRANSFORM, -1, 2, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_MODULAR_COLOR_SPACE, -1, 41, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_MODULAR_GROUP_SIZE, -1, 3, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_MODULAR_PREDICTOR, -1, 15, 0},

      IntOptionSpec{JXL_ENC_FRAME_SETTING_MODULAR_NB_PREV_CHANNELS, -1, 11, 0},

  };

  struct FloatOptionSpec {

    JxlEncoderFrameSettingId flag;

    float possible_values[4];

    float value;

  };

  std::vector<FloatOptionSpec> float_options = {

      FloatOptionSpec{JXL_ENC_FRAME_SETTING_CHANNEL_COLORS_GLOBAL_PERCENT,

                      {-1, 0, 50, 100},

                      -1},

      FloatOptionSpec{JXL_ENC_FRAME_SETTING_CHANNEL_COLORS_GROUP_PERCENT,

                      {-1, 0, 50, 100},

                      -1},

      FloatOptionSpec{JXL_ENC_FRAME_SETTING_MODULAR_MA_TREE_LEARNING_PERCENT,

                      {-1, 0, 50, 100},

                      -1},

      FloatOptionSpec{

          JXL_ENC_FRAME_SETTING_PHOTON_NOISE, {-1, 200, 1600, 10000}, -1},

  };

  uint8_t num_threads;

  float distance;  // 0.01 - 25

  // Tiled to cover the entire image area.

  uint16_t pixel_data[4][64][64];

  static FuzzSpec FromData(const uint8_t* data, size_t len) {

    size_t pos = 0;

    auto u8 = [&]() -> uint8_t {

      if (pos == len) return 0;

      return data[pos++];

    };

    auto b1 = [&]() -> bool { return static_cast<bool>(u8() % 2); };

    auto u16 = [&]() -> uint16_t { return (uint16_t{u8()} << 8) | u8(); };

    FuzzSpec spec;

    spec.xsize = uint32_t{u16()} + 1;

    spec.ysize = uint32_t{u16()} + 1;

    constexpr uint64_t kMaxSize = 1 << 24;

    if (spec.xsize * uint64_t{spec.ysize} > kMaxSize) {

      spec.ysize = kMaxSize / spec.xsize;

    }

    spec.grayscale = b1();

    spec.alpha = b1();

    spec.bit_depth = u8() % 16 + 1;

    // constants chosen so to cover the entire 0.01 - 25 range.

    spec.distance = u8() % 2 ? 0.0 : 0.01 + 0.00038132 * u16();

    Check(spec.float_options[2].flag ==

          JXL_ENC_FRAME_SETTING_MODULAR_MA_TREE_LEARNING_PERCENT);

    Check(spec.int_options[15].flag == JXL_ENC_FRAME_SETTING_COLOR_TRANSFORM);

    if (spec.distance != 0 || spec.int_options[15].value == 0) {

      spec.float_options[2].possible_values[1] = 1;

    }

    spec.num_threads = u8();

    for (auto& int_opt : spec.int_options) {

      int_opt.value = u8() % (int_opt.max - int_opt.min + 1) + int_opt.min;

    }

    for (auto& float_opt : spec.float_options) {

      float_opt.value = float_opt.possible_values[u8() % 4];

    }

    for (auto& x : spec.pixel_data) {

      for (auto& y : x) {

        for (auto& p : y) {

          p = u16();

        }

      }

    }

    return spec;

  }

};

StatusOr<std::vector<uint8_t>> Encode(const FuzzSpec& spec,

                                      TrackingMemoryManager& memory_manager,

                                      bool streaming) {

  auto runner = JxlThreadParallelRunnerMake(nullptr, spec.num_threads);

  JxlEncoderPtr enc_ptr = JxlEncoderMake(memory_manager.get());

  JxlEncoder* enc = enc_ptr.get();

  Check(JxlEncoderSetParallelRunner(enc, JxlThreadParallelRunner,

                                    runner.get()) == JXL_ENC_SUCCESS);

  JxlEncoderFrameSettings* frame_settings =

      JxlEncoderFrameSettingsCreate(enc, nullptr);

  Check(JxlEncoderSetFrameDistance(frame_settings, spec.distance) ==

        JXL_ENC_SUCCESS);

  for (const auto& opt : spec.int_options) {

    Check(JxlEncoderFrameSettingsSetOption(frame_settings, opt.flag,

                                           opt.value) == JXL_ENC_SUCCESS);

  }

  for (const auto& opt : spec.float_options) {

    if (opt.value != -1) {

      Check(JxlEncoderFrameSettingsSetFloatOption(

                frame_settings, opt.flag, opt.value) == JXL_ENC_SUCCESS);

    }

  }

  Check(JxlEncoderFrameSettingsSetOption(frame_settings,

                                         JXL_ENC_FRAME_SETTING_BUFFERING,

                                         streaming ? 3 : 0) == JXL_ENC_SUCCESS);

  JxlBasicInfo basic_info;

  JxlEncoderInitBasicInfo(&basic_info);

  basic_info.num_color_channels = spec.grayscale ? 1 : 3;

  basic_info.xsize = spec.xsize;

  basic_info.ysize = spec.ysize;

  basic_info.bits_per_sample = spec.bit_depth;

  basic_info.uses_original_profile = JXL_FALSE;

  uint32_t nchan = basic_info.num_color_channels;

  if (spec.alpha) {

    nchan += 1;

    basic_info.alpha_bits = spec.bit_depth;

    basic_info.num_extra_channels = 1;

  }

  Check(JxlEncoderSetBasicInfo(enc, &basic_info) == JXL_ENC_SUCCESS);

  if (spec.alpha) {

    JxlExtraChannelInfo info;

    memset(&info, 0, sizeof(info));

    info.type = JxlExtraChannelType::JXL_CHANNEL_ALPHA;

    info.bits_per_sample = spec.bit_depth;

    JxlEncoderSetExtraChannelInfo(enc, 0, &info);

  }

  JxlColorEncoding color_encoding;

  memset(&color_encoding, 0, sizeof(color_encoding));

  color_encoding.color_space = spec.grayscale

                                   ? JxlColorSpace::JXL_COLOR_SPACE_GRAY

                                   : JxlColorSpace::JXL_COLOR_SPACE_RGB;

  color_encoding.transfer_function =

      JxlTransferFunction::JXL_TRANSFER_FUNCTION_SRGB;

  color_encoding.primaries = JxlPrimaries::JXL_PRIMARIES_2100;

  color_encoding.white_point = JxlWhitePoint::JXL_WHITE_POINT_D65;

  color_encoding.rendering_intent =

      JxlRenderingIntent::JXL_RENDERING_INTENT_RELATIVE;

  Check(JxlEncoderSetColorEncoding(enc, &color_encoding) == JXL_ENC_SUCCESS);

  JxlFrameHeader frame_header;

  JxlEncoderInitFrameHeader(&frame_header);

  // TODO(szabadka) Add more frame header options.

  Check(JxlEncoderSetFrameHeader(frame_settings, &frame_header) ==

        JXL_ENC_SUCCESS);

  JxlPixelFormat pixelformat = {nchan, JXL_TYPE_UINT16, JXL_LITTLE_ENDIAN, 0};

  std::vector<uint16_t> pixels(spec.xsize * static_cast<uint64_t>(spec.ysize) *

                               nchan);

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

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

      for (size_t c = 0; c < nchan; c++) {

        pixels[(y * spec.xsize + x) * nchan + c] =

            spec.pixel_data[c][y % 64][x % 64];

      }

    }

  }

  JxlEncoderStatus status =

      JxlEncoderAddImageFrame(frame_settings, &pixelformat, pixels.data(),

                              pixels.size() * sizeof(uint16_t));

  // TODO(eustas): update when API will provide OOM status.

  if (memory_manager.seen_oom) {

    // Actually, that is fine.

    return JXL_FAILURE("OOM");

  }

  Check(status == JXL_ENC_SUCCESS);

  JxlEncoderCloseInput(enc);

  // Reading compressed output

  JxlEncoderStatus process_result = JXL_ENC_NEED_MORE_OUTPUT;

  std::vector<uint8_t> buf(1024);

  size_t written = 0;

  while (process_result == JXL_ENC_NEED_MORE_OUTPUT) {

    buf.resize(buf.size() * 2);

    uint8_t* next_out = buf.data() + written;

    size_t avail_out = buf.size() - written;

    process_result = JxlEncoderProcessOutput(enc, &next_out, &avail_out);

    written = next_out - buf.data();

  }

  // TODO(eustas): update when API will provide OOM status.

  if (memory_manager.seen_oom) {

    // Actually, that is fine.

    return JXL_FAILURE("OOM");

  }

  Check(process_result == JXL_ENC_SUCCESS);

  buf.resize(written);

  return buf;

}

StatusOr<std::vector<float>> Decode(const std::vector<uint8_t>& data,

                                    TrackingMemoryManager& memory_manager) {

  // Multi-threaded parallel runner.

  auto dec = JxlDecoderMake(memory_manager.get());

  Check(JxlDecoderSubscribeEvents(dec.get(),

                                  JXL_DEC_BASIC_INFO | JXL_DEC_FULL_IMAGE) ==

        JXL_DEC_SUCCESS);

  JxlBasicInfo info;

  JxlPixelFormat format = {3, JXL_TYPE_FLOAT, JXL_NATIVE_ENDIAN, 0};

  std::vector<float> pixels;

  JxlDecoderSetInput(dec.get(), data.data(), data.size());

  JxlDecoderCloseInput(dec.get());

  for (;;) {

    JxlDecoderStatus status = JxlDecoderProcessInput(dec.get());

    if (status == JXL_DEC_BASIC_INFO) {

      Check(JxlDecoderGetBasicInfo(dec.get(), &info) == JXL_DEC_SUCCESS);

    } else if (status == JXL_DEC_NEED_IMAGE_OUT_BUFFER) {

      size_t buffer_size;

      Check(JxlDecoderImageOutBufferSize(dec.get(), &format, &buffer_size) ==

            JXL_DEC_SUCCESS);

      pixels.resize(buffer_size / sizeof(float));

      void* pixels_buffer = static_cast<void*>(pixels.data());

      size_t pixels_buffer_size = pixels.size() * sizeof(float);

      Check(JxlDecoderSetImageOutBuffer(dec.get(), &format, pixels_buffer,

                                        pixels_buffer_size) == JXL_DEC_SUCCESS);

    } else if (status == JXL_DEC_FULL_IMAGE || status == JXL_DEC_SUCCESS) {

      return pixels;

    } else {

      // TODO(eustas): update when API will provide OOM status.

      if (memory_manager.seen_oom) {

        // Actually, that is fine.

        return JXL_FAILURE("OOM");

      }

      // Unexpected status

      Check(false);

    }

  }

}

Status Run(const FuzzSpec& spec, TrackingMemoryManager& memory_manager) {

  std::vector<uint8_t> enc_default;

  std::vector<uint8_t> enc_streaming;

  const auto encode = [&]() -> Status {

    // It is not clear, which approach eatc more memory.

    JXL_ASSIGN_OR_RETURN(enc_default, Encode(spec, memory_manager, false));

    Check(memory_manager.Reset());

    JXL_ASSIGN_OR_RETURN(enc_streaming, Encode(spec, memory_manager, true));

    Check(memory_manager.Reset());

    return true;

  };

  // It is fine, if encoder OOMs.

  if (!encode()) return true;

  // It is NOT OK, it decoder OOMs - it should not consume more than encoder.

  JXL_ASSIGN_OR_RETURN(auto dec_default, Decode(enc_default, memory_manager));

  Check(memory_manager.Reset());

  JXL_ASSIGN_OR_RETURN(auto dec_streaming,

                       Decode(enc_streaming, memory_manager));

  Check(memory_manager.Reset());

  Check(dec_default == dec_streaming);

  return true;

}

int DoTestOneInput(const uint8_t* data, size_t size) {

  auto spec = FuzzSpec::FromData(data, size);

  TrackingMemoryManager memory_manager{/* cap */ 1 * kGiB,

                                       /* total_cap */ 5 * kGiB};

  Check(Run(spec, memory_manager));

  return 0;

}

}  // namespace

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size) {

  return DoTestOneInput(data, size);

}

void TestOneInput(const std::vector<uint8_t>& data) {

  DoTestOneInput(data.data(), data.size());

}

FUZZ_TEST(StreamingFuzzTest, TestOneInput);