// 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 <array>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <utility>
#include <vector>

#include "lib/jxl/fuzztest.h"
#include "tools/tracking_memory_manager.h"

namespace {

using ::jpegxl::tools::kGiB;
using ::jpegxl::tools::TrackingMemoryManager;

constexpr size_t kMemoryCap = kGiB;  // enough for 85.3MPx without overhead
constexpr size_t kBaseMaxSize = 16 << 20;  // 16MPx

void CheckImpl(bool ok, const char* condition, const char* file, int line) {
  if (!ok) {
    fprintf(stderr, "Check(%s) failed at %s:%d\n", condition, file, line);
    __builtin_trap();
  }
}
#define Check(OK) CheckImpl((OK), #OK, __FILE__, __LINE__)

float Clamp1(float val, float low, float hi) {
  return val < low ? low : val > hi ? hi : val;
}

size_t RoundUpTo(size_t what, size_t align) {
  size_t reminder = what % align;
  return (reminder == 0) ? what : (what + align - reminder);
}

struct FuzzSpec {
  uint32_t xsize;
  uint32_t ysize;
  bool grayscale;
  bool alpha;
  uint8_t bit_depth;  // 1 - 16

  struct IntOptionSpec {
    JxlEncoderFrameSettingId flag;
    std::string name;
    int min;
    int max;
    int value;
  };

#define INT_OPTION(FLAG, MIN_V, MAX_V, V) \
  IntOptionSpec { FLAG, #FLAG, MIN_V, MAX_V, V }

  std::vector<IntOptionSpec> int_options = {
      INT_OPTION(JXL_ENC_FRAME_SETTING_EFFORT, 1, 9, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_DECODING_SPEED, 0, 4, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_NOISE, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_DOTS, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_PATCHES, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_EPF, -1, 3, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_GABORISH, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_MODULAR, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_KEEP_INVISIBLE, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_RESPONSIVE, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_PROGRESSIVE_AC, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_QPROGRESSIVE_AC, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_PROGRESSIVE_DC, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_PALETTE_COLORS, -1, 255, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_LOSSY_PALETTE, -1, 1, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_COLOR_TRANSFORM, -1, 2, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_MODULAR_COLOR_SPACE, -1, 41, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_MODULAR_GROUP_SIZE, -1, 3, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_MODULAR_PREDICTOR, -1, 15, 0),
      INT_OPTION(JXL_ENC_FRAME_SETTING_MODULAR_NB_PREV_CHANNELS, -1, 11, 0),
  };

#undef INT_OPTION

  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;
    // TODO(eustas): allow dimensions to be 130k
    spec.xsize = uint32_t{u16()} + 1;
    spec.ysize = uint32_t{u16()} + 1;
    spec.grayscale = b1();
    spec.alpha = b1();
    spec.bit_depth = u8() % 16 + 1;
    // constants chosen so to cover the entire 0.01 - 25 range.
    bool lossless = ((u8() % 2) == 1);
    spec.distance = lossless ? 0.0 : 0.01 + 0.00038132 * u16();

    spec.num_threads = u8() & 0xF;

    for (auto& int_opt : spec.int_options) {
      int_opt.value = u8() % (int_opt.max - int_opt.min + 1) + int_opt.min;
    }

    Check(spec.int_options[15].flag == JXL_ENC_FRAME_SETTING_COLOR_TRANSFORM);
    if (!lossless || spec.int_options[15].value == 0) {
      Check(spec.float_options[2].flag ==
            JXL_ENC_FRAME_SETTING_MODULAR_MA_TREE_LEARNING_PERCENT);
      spec.float_options[2].possible_values[1] = 1;
    }

    for (auto& float_opt : spec.float_options) {
      float_opt.value = float_opt.possible_values[u8() % 4];
    }

    Check(spec.int_options[7].flag == JXL_ENC_FRAME_SETTING_MODULAR);
    bool modular = (spec.int_options[7].value == 1);
    Check(spec.int_options[18].flag == JXL_ENC_FRAME_SETTING_MODULAR_PREDICTOR);
    bool slow_predictor = (spec.int_options[18].value >= 14);
    uint64_t max_size = kBaseMaxSize;
    if (modular && slow_predictor) max_size /= 2;
    if (sizeof(size_t) == 4) max_size /= 1.5;
    constexpr size_t group_dim = 256;
    uint64_t in_mem_xsize = RoundUpTo(spec.xsize, group_dim);
    if (in_mem_xsize * spec.ysize > max_size) {
      spec.ysize = max_size / in_mem_xsize;
      Check(spec.ysize > 0);
    }
    uint64_t in_mem_ysize = RoundUpTo(spec.ysize, group_dim);
    if (spec.xsize * in_mem_ysize > max_size) {
      spec.xsize = max_size / in_mem_ysize;
      Check(spec.xsize > 0);
    }

    for (auto& x : spec.pixel_data) {
      for (auto& y : x) {
        for (auto& p : y) {
          p = u16();
        }
      }
    }

    if (false) {
      fprintf(stderr, "Image size: %d X %d, d=%f, num_threads: %d\n",
              spec.xsize, spec.ysize, spec.distance, spec.num_threads);
      for (auto& int_opt : spec.int_options) {
        fprintf(stderr, "%s = %d\n", int_opt.name.c_str(), int_opt.value);
      }
    }

    return spec;
  }
};

bool Encode(const FuzzSpec& spec, TrackingMemoryManager& memory_manager,
            bool streaming, std::vector<uint8_t>& result) {
  result.clear();
  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(frame_settings != 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;
  bool lossless = (spec.distance == 0.0f);
  basic_info.uses_original_profile = TO_JXL_BOOL(lossless);
  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++) {
        // TODO(eustas): make it less regular for tiles except (0, 0)
        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 false;
  }
  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 false;
  }
  Check(process_result == JXL_ENC_SUCCESS);
  buf.resize(written);

  result = std::move(buf);
  return true;
}

bool Decode(const std::vector<uint8_t>& data,
            TrackingMemoryManager& memory_manager, std::vector<float>& pixels) {
  // 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};

  pixels.clear();

  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 true;
    } else {
      // TODO(eustas): update when API will provide OOM status.
      if (memory_manager.seen_oom) {
        // Actually, that is fine.
        return false;
      }
      // Unexpected status
      Check(false);
    }
  }
}

void Run(const FuzzSpec& spec) {
  size_t memory_cap = kMemoryCap;
  size_t total_cap_multiplier = 5;
  if (spec.xsize < 64 || spec.ysize < 64) {
    total_cap_multiplier = 20;
  }
  TrackingMemoryManager memory_manager{memory_cap,
                                       memory_cap * total_cap_multiplier};

  std::vector<uint8_t> enc_default;
  std::vector<uint8_t> enc_streaming;

  const auto encode = [&]() -> bool {
    // It is not clear, which approach eats more memory.
    if (!Encode(spec, memory_manager, false, enc_default)) return false;
    Check(memory_manager.Reset());
    if (!Encode(spec, memory_manager, true, enc_streaming)) return false;
    Check(memory_manager.Reset());
    return true;
  };
  // It is fine, if encoder OOMs.
  if (!encode()) return;

  // It is NOT OK, if decoder OOMs - it should not consume more than encoder.
  std::vector<float> dec_default;
  Check(Decode(enc_default, memory_manager, dec_default));
  Check(memory_manager.Reset());
  std::vector<float> dec_streaming;
  Check(Decode(enc_streaming, memory_manager, dec_streaming));
  Check(memory_manager.Reset());

  Check(dec_default.size() == dec_streaming.size());

  Check(spec.int_options[0].flag == JXL_ENC_FRAME_SETTING_EFFORT);
  int effort = spec.int_options[0].value;
  std::array<float, 10> kThreshold = {0.00f,   0.05f,   0.05f,   0.05f, 0.05f,
                                      0.0625f, 0.0625f, 0.0625f, 0.10f, 0.10f};
  float threshold = kThreshold[effort];

  int outlier_count = 0;
  for (size_t i = 0; i < dec_default.size(); ++i) {
    float d1 = Clamp1(dec_default[i], 0.0f, 1.0f);
    float d2 = Clamp1(dec_streaming[i], 0.0f, 1.0f);
    float abs_diff = std::abs(d1 - d2);
    if (abs_diff > threshold) outlier_count++;
  }
  if (false) {
    fprintf(stderr, "Number of outlier values: %d / %d\n", outlier_count,
            static_cast<int>(dec_default.size()));
  }
  Check(outlier_count == 0);
}

int DoTestOneInput(const uint8_t* data, size_t size) {
  auto spec = FuzzSpec::FromData(data, size);

  Run(spec);
  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);
