/*
 * HEIF codec.
 * Copyright (c) 2019 struktur AG, Joachim Bauch <bauch@struktur.de>
 *
 * This file is part of libheif.
 *
 * libheif is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as
 * published by the Free Software Foundation, either version 3 of
 * the License, or (at your option) any later version.
 *
 * libheif is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with libheif.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <assert.h>

#include <sstream>

#include "bitstream.h"
#include "color-conversion/colorconversion.h"
#include "pixelimage.h"

static bool is_valid_chroma(uint8_t chroma)
{
  switch (chroma) {
    case heif_chroma_monochrome:
    case heif_chroma_420:
    case heif_chroma_422:
    case heif_chroma_444:
    case heif_chroma_interleaved_RGB:
    case heif_chroma_interleaved_RGBA:
    case heif_chroma_interleaved_RRGGBB_BE:
    case heif_chroma_interleaved_RRGGBBAA_BE:
    case heif_chroma_interleaved_RRGGBB_LE:
    case heif_chroma_interleaved_RRGGBBAA_LE:
      return true;
    default:
      return false;
  }
}

static bool is_valid_colorspace(uint8_t colorspace)
{
  switch (colorspace) {
    case heif_colorspace_YCbCr:
    case heif_colorspace_RGB:
    case heif_colorspace_monochrome:
      return true;
    default:
      return false;
  }
}

static bool read_plane(BitstreamRange* range,
                       std::shared_ptr<HeifPixelImage> image, heif_channel channel,
                       uint32_t width, uint32_t height, int bit_depth)
{
  if (width <= 0 || height <= 0) {
    return false;
  }
  if (std::numeric_limits<size_t>::max()/width/height == 0) {
    return false;
  }
  if (!range->prepare_read(static_cast<size_t>(width) * height)) {
    return false;
  }
  if (auto err = image->add_plane(channel, width, height, bit_depth, heif_get_disabled_security_limits())) {
    return false;
  }
  size_t stride;
  uint8_t* plane = image->get_plane(channel, &stride);
  assert(stride >= width);
  auto stream = range->get_istream();
  for (uint32_t y = 0; y < height; y++, plane += stride) {
    assert(stream->read(plane, width));
  }
  return true;
}

static bool read_plane_interleaved(BitstreamRange* range,
                                   std::shared_ptr<HeifPixelImage> image, heif_channel channel,
                                   uint32_t width, uint32_t height, int bit_depth, int comps)
{
  if (width <= 0 || height <= 0) {
    return false;
  }
  if (std::numeric_limits<size_t>::max()/width/height/comps == 0) {
    return false;
  }
  if (!range->prepare_read(static_cast<size_t>(width) * height * comps)) {
    return false;
  }
  if (auto err = image->add_plane(channel, width, height, bit_depth, heif_get_disabled_security_limits())) {
    return false;
  }
  size_t stride;
  uint8_t* plane = image->get_plane(channel, &stride);
  assert(stride >= width * comps);
  auto stream = range->get_istream();
  for (uint32_t y = 0; y < height; y++, plane += stride) {
    assert(stream->read(plane, width * comps));
  }
  return true;
}

extern "C" int LLVMFuzzerTestOneInput(const uint8_t* data, size_t size)
{
  auto reader = std::make_shared<StreamReader_memory>(data, size, false);
  BitstreamRange range(reader, size);

  uint32_t width;
  uint32_t height;
  int bit_depth;
  bool alpha;
  uint8_t in_chroma;
  uint8_t in_colorspace;
  uint8_t out_chroma;
  uint8_t out_colorspace;
  if (!range.prepare_read(10)) {
    return 0;
  }

  width = range.read16();
  height = range.read16();
  bit_depth = range.read8();
  alpha = range.read8() == 1;
  in_chroma = range.read8();
  in_colorspace = range.read8();
  out_chroma = range.read8();
  out_colorspace = range.read8();

  // Width / height must be a multiple of 2.
  if (width == 0 || height == 0 || (width & 1) != 0 || (height & 1) != 0) {
    return 0;
  }

  switch (bit_depth) {
    case 8:
      break;
    default:
      // TODO: Add support for more color depths.
      return 0;
  }

  if (!is_valid_chroma(in_chroma) || !is_valid_colorspace(in_colorspace) ||
      !is_valid_chroma(out_chroma) || !is_valid_colorspace(out_colorspace)) {
    return 0;
  }

  auto in_image = std::make_shared<HeifPixelImage>();
  in_image->create(width, height, static_cast<heif_colorspace>(in_colorspace),
                   static_cast<heif_chroma>(in_chroma));

  switch (in_colorspace) {
    case heif_colorspace_YCbCr:
      switch (in_chroma) {
        case heif_chroma_420:
          if (!read_plane(&range, in_image, heif_channel_Y,
                          width, height, bit_depth)) {
            return 0;
          }
          if (!read_plane(&range, in_image, heif_channel_Cb,
                          width / 2, height / 2, bit_depth)) {
            return 0;
          }
          if (!read_plane(&range, in_image, heif_channel_Cr,
                          width / 2, height / 2, bit_depth)) {
            return 0;
          }
          break;
        case heif_chroma_422:
          if (!read_plane(&range, in_image, heif_channel_Y,
                          width, height, bit_depth)) {
            return 0;
          }
          if (!read_plane(&range, in_image, heif_channel_Cb,
                          width / 2, height, bit_depth)) {
            return 0;
          }
          if (!read_plane(&range, in_image, heif_channel_Cr,
                          width / 2, height, bit_depth)) {
            return 0;
          }
          break;
        case heif_chroma_444:
          if (!read_plane(&range, in_image, heif_channel_Y,
                          width, height, bit_depth)) {
            return 0;
          }
          if (!read_plane(&range, in_image, heif_channel_Cb,
                          width, height, bit_depth)) {
            return 0;
          }
          if (!read_plane(&range, in_image, heif_channel_Cr,
                          width, height, bit_depth)) {
            return 0;
          }
          break;
        default:
          return 0;
      }
      break;
    case heif_colorspace_RGB:
      switch (in_chroma) {
        case heif_chroma_interleaved_RGB:
          if (!read_plane_interleaved(&range, in_image,
                                      heif_channel_interleaved, width, height, bit_depth, 3)) {
            return 0;
          }
          break;
        case heif_chroma_interleaved_RGBA:
          if (!read_plane_interleaved(&range, in_image,
                                      heif_channel_interleaved, width, height, bit_depth, 4)) {
            return 0;
          }
          alpha = false;  // Already part of interleaved data.
          break;
        default:
          // TODO: Support other RGB chromas.
          return 0;
      }
      break;
    case heif_colorspace_monochrome:
      if (in_chroma != heif_chroma_monochrome) {
        return 0;
      }
      if (!read_plane(&range, in_image, heif_channel_Y,
                      width, height, bit_depth)) {
        return 0;
      }
      break;
    default:
      assert(false);
  }

  if (alpha) {
    if (!read_plane(&range, in_image, heif_channel_Alpha,
                    width, height, bit_depth)) {
      return 0;
    }
  }

  // TODO: also fuzz these parameters.
  int output_bpp = 0; // Same as input.
  heif_encoding_options* options = heif_encoding_options_alloc();

  auto out_image_result = convert_colorspace(in_image,
                                             static_cast<heif_colorspace>(out_colorspace),
                                             static_cast<heif_chroma>(out_chroma),
                                             nullptr,
                                             output_bpp,
                                             options->color_conversion_options,
                                             heif_get_disabled_security_limits());

  heif_encoding_options_free(options);

  if (out_image_result.error) {
    // Conversion is not supported.
    return 0;
  }

  auto out_image = *out_image_result;

  assert(out_image->get_width() == width);
  assert(out_image->get_height() == height);
  assert(out_image->get_chroma_format() ==
         static_cast<heif_chroma>(out_chroma));
  assert(out_image->get_colorspace() ==
         static_cast<heif_colorspace>(out_colorspace));
  return 0;
}