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

#include <algorithm>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <hwy/targets.h>
#include <vector>

#include "lib/jpegli/common.h"
#include "lib/jpegli/decode.h"
#include "lib/jpegli/fuzztest.h"
#include "lib/jpegli/types.h"

namespace {

// Externally visible value to ensure pixels are used in the fuzzer.
int external_code = 0;

template <typename It>
void Consume(const It& begin, const It& end) {
  for (auto it = begin; it < end; ++it) {
    if (*it == 0) {
      external_code ^= ~0;
    } else {
      external_code ^= *it;
    }
  }
}

// Options for the fuzzing
struct FuzzSpec {
  size_t chunk_size;
  JpegliDataType output_type;
  JpegliEndianness output_endianness;
  int crop_output;
};

constexpr uint8_t kFakeEoiMarker[2] = {0xff, 0xd9};
constexpr size_t kNumSourceBuffers = 4;

class SourceManager {
 public:
  SourceManager(const uint8_t* data, size_t len, size_t max_chunk_size)
      : data_(data), len_(len), max_chunk_size_(max_chunk_size) {
    pub_.skip_input_data = skip_input_data;
    pub_.resync_to_restart = jpegli_resync_to_restart;
    pub_.term_source = term_source;
    pub_.init_source = init_source;
    pub_.fill_input_buffer = fill_input_buffer;
    if (max_chunk_size_ == 0) max_chunk_size_ = len;
    buffers_.resize(kNumSourceBuffers, std::vector<uint8_t>(max_chunk_size_));
    Reset();
  }

  void Reset() {
    pub_.next_input_byte = nullptr;
    pub_.bytes_in_buffer = 0;
    pos_ = 0;
    chunk_idx_ = 0;
  }

 private:
  jpeg_source_mgr pub_;
  const uint8_t* data_;
  size_t len_;
  size_t chunk_idx_;
  size_t pos_;
  size_t max_chunk_size_;
  std::vector<std::vector<uint8_t>> buffers_;

  static void init_source(j_decompress_ptr cinfo) {}

  static boolean fill_input_buffer(j_decompress_ptr cinfo) {
    auto* src = reinterpret_cast<SourceManager*>(cinfo->src);
    if (src->pos_ < src->len_) {
      size_t remaining = src->len_ - src->pos_;
      size_t chunk_size = std::min(remaining, src->max_chunk_size_);
      size_t next_idx = ++src->chunk_idx_ % kNumSourceBuffers;
      // Larger number of chunks causes fuzzer timuout.
      if (src->chunk_idx_ >= (1u << 15)) {
        chunk_size = remaining;
        next_idx = src->buffers_.size();
        src->buffers_.emplace_back(chunk_size);
      }
      uint8_t* next_buffer = src->buffers_[next_idx].data();
      memcpy(next_buffer, src->data_ + src->pos_, chunk_size);
      src->pub_.next_input_byte = next_buffer;
      src->pub_.bytes_in_buffer = chunk_size;
    } else {
      src->pub_.next_input_byte = kFakeEoiMarker;
      src->pub_.bytes_in_buffer = 2;
      src->len_ += 2;
    }
    src->pos_ += src->pub_.bytes_in_buffer;
    return TRUE;
  }

  static void skip_input_data(j_decompress_ptr cinfo,
                              long num_bytes /* NOLINT */) {
    auto* src = reinterpret_cast<SourceManager*>(cinfo->src);
    if (num_bytes <= 0) {
      return;
    }
    if (src->pub_.bytes_in_buffer >= static_cast<size_t>(num_bytes)) {
      src->pub_.bytes_in_buffer -= num_bytes;
      src->pub_.next_input_byte += num_bytes;
    } else {
      src->pos_ += num_bytes - src->pub_.bytes_in_buffer;
      src->pub_.bytes_in_buffer = 0;
    }
  }

  static void term_source(j_decompress_ptr cinfo) {}
};

bool DecodeJpeg(const uint8_t* data, size_t size, size_t max_pixels,
                const FuzzSpec& spec, std::vector<uint8_t>* pixels,
                size_t* xsize, size_t* ysize) {
  SourceManager src(data, size, spec.chunk_size);
  jpeg_decompress_struct cinfo;
  const auto try_catch_block = [&]() -> bool {
    jpeg_error_mgr jerr;
    jmp_buf env;
    cinfo.err = jpegli_std_error(&jerr);
    if (setjmp(env)) {
      return false;
    }
    cinfo.client_data = reinterpret_cast<void*>(&env);
    cinfo.err->error_exit = [](j_common_ptr cinfo) {
      jmp_buf* env = reinterpret_cast<jmp_buf*>(cinfo->client_data);
      jpegli_destroy(cinfo);
      longjmp(*env, 1);
    };
    cinfo.err->emit_message = [](j_common_ptr cinfo, int msg_level) {};
    jpegli_create_decompress(&cinfo);
    cinfo.src = reinterpret_cast<jpeg_source_mgr*>(&src);
    jpegli_read_header(&cinfo, TRUE);
    *xsize = cinfo.image_width;
    *ysize = cinfo.image_height;
    size_t num_pixels = *xsize * *ysize;
    if (num_pixels > max_pixels) return false;
    jpegli_set_output_format(&cinfo, spec.output_type, spec.output_endianness);
    jpegli_start_decompress(&cinfo);
    if (spec.crop_output) {
      JDIMENSION xoffset = cinfo.output_width / 3;
      JDIMENSION xsize_cropped = cinfo.output_width / 3;
      jpegli_crop_scanline(&cinfo, &xoffset, &xsize_cropped);
    }

    size_t bytes_per_sample = jpegli_bytes_per_sample(spec.output_type);
    size_t stride =
        bytes_per_sample * cinfo.output_components * cinfo.output_width;
    size_t buffer_size = *ysize * stride;
    pixels->resize(buffer_size);
    for (size_t y = 0; y < *ysize; ++y) {
      JSAMPROW rows[] = {pixels->data() + y * stride};
      jpegli_read_scanlines(&cinfo, rows, 1);
    }
    Consume(pixels->cbegin(), pixels->cend());
    jpegli_finish_decompress(&cinfo);
    return true;
  };
  bool success = try_catch_block();
  jpegli_destroy_decompress(&cinfo);
  return success;
}

int DoTestOneInput(const uint8_t* data, size_t size) {
  if (size < 4) return 0;
  uint32_t flags = 0;
  size_t used_flag_bits = 0;
  memcpy(&flags, data + size - 4, 4);
  size -= 4;

  const auto getFlag = [&flags, &used_flag_bits](size_t max_value) {
    size_t limit = 1;
    while (limit <= max_value) {
      limit <<= 1;
      used_flag_bits++;
      if (used_flag_bits > 32) abort();
    }
    uint32_t result = flags % limit;
    flags /= limit;
    return result % (max_value + 1);
  };

  FuzzSpec spec;
  spec.output_type = static_cast<JpegliDataType>(getFlag(JPEGLI_TYPE_UINT16));
  spec.output_endianness =
      static_cast<JpegliEndianness>(getFlag(JPEGLI_BIG_ENDIAN));
  uint32_t chunks = getFlag(15);
  spec.chunk_size = chunks ? 1u << (chunks - 1) : 0;
  spec.crop_output = getFlag(1);

  std::vector<uint8_t> pixels;
  size_t xsize;
  size_t ysize;
  size_t max_pixels = 1 << 21;

  const auto targets = hwy::SupportedAndGeneratedTargets();
  hwy::SetSupportedTargetsForTest(targets[getFlag(targets.size() - 1)]);
  DecodeJpeg(data, size, max_pixels, spec, &pixels, &xsize, &ysize);
  hwy::SetSupportedTargetsForTest(0);

  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(JpegliDecFuzzTest, TestOneInput);