/**
 * File: HEIF IO
 *
 * Read and write HEIF images.
 */

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif /* HAVE_CONFIG_H */


#include <stdio.h>
#include <math.h>
#include <string.h>
#include <stdlib.h>
#include "gd.h"
#include "gd_errors.h"
#include "gdhelpers.h"

#ifdef HAVE_LIBHEIF
#include <libheif/heif.h>

#define GD_HEIF_ALLOC_STEP (4*1024)
#define GD_HEIF_HEADER 12

typedef enum gd_heif_brand {
	GD_HEIF_BRAND_AVIF = 1,
	GD_HEIF_BRAND_MIF1 = 2,
	GD_HEIF_BRAND_HEIC = 4,
	GD_HEIF_BRAND_HEIX = 8,
} gd_heif_brand;

/*
  Function: gdImageCreateFromHeif

    <gdImageCreateFromHeif> is called to load truecolor images from
    HEIF format files. Invoke <gdImageCreateFromHeif> with an
    already opened pointer to a file containing the desired
    image. <gdImageCreateFromHeif> returns a <gdImagePtr> to the new
    truecolor image, or NULL if unable to load the image (most often
    because the file is corrupt or does not contain a HEIF
    image). <gdImageCreateFromHeif> does not close the file.

    You can inspect the sx and sy members of the image to determine
    its size. The image must eventually be destroyed using
    <gdImageDestroy>.

    *The returned image is always a truecolor image.*

  Parameters:

    infile - The input FILE pointer.

  Returns:

    A pointer to the new *truecolor* image.  This will need to be
    destroyed with <gdImageDestroy> once it is no longer needed.

    On error, returns NULL.
*/
BGD_DECLARE(gdImagePtr) gdImageCreateFromHeif(FILE *inFile)
{
	gdImagePtr im;
	gdIOCtx *in = gdNewFileCtx(inFile);

	if (!in)
		return NULL;
	im = gdImageCreateFromHeifCtx(in);
	in->gd_free(in);

	return im;
}

/*
  Function: gdImageCreateFromHeifPtr

    See <gdImageCreateFromHeif>.

  Parameters:

    size            - size of HEIF data in bytes.
    data            - pointer to HEIF data.
*/
BGD_DECLARE(gdImagePtr) gdImageCreateFromHeifPtr(int size, void *data)
{
	gdImagePtr im;
	gdIOCtx *in = gdNewDynamicCtxEx(size, data, 0);

	if (!in)
		return NULL;
	im = gdImageCreateFromHeifCtx(in);
	in->gd_free(in);

	return im;
}

static int _gdHeifCheckBrand(unsigned char *magic, gd_heif_brand expected_brand)
{
	if (memcmp(magic + 4, "ftyp", 4) != 0)
		return GD_FALSE;
	if (memcmp(magic + 8, "avif", 4) == 0 && expected_brand & GD_HEIF_BRAND_AVIF)
		return GD_TRUE;
	if (memcmp(magic + 8, "heic", 4) == 0 && expected_brand & GD_HEIF_BRAND_HEIC)
		return GD_TRUE;
	if (memcmp(magic + 8, "heix", 4) == 0 && expected_brand & GD_HEIF_BRAND_HEIX)
		return GD_TRUE;
	if (memcmp(magic + 8, "mif1", 4) == 0 && expected_brand & GD_HEIF_BRAND_MIF1)
		return GD_TRUE;

	return GD_FALSE;
}

static gdImagePtr _gdImageCreateFromHeifCtx(gdIOCtx *infile, gd_heif_brand expected_brand)
{
	struct heif_context *heif_ctx;
	struct heif_decoding_options *heif_dec_opts;
	struct heif_image_handle *heif_imhandle;
	struct heif_image *heif_im;
	struct heif_error err;
	int    width, height;
	uint8_t   *filedata = NULL;
	uint8_t    *rgba = NULL;
	unsigned char   *read, *temp, magic[GD_HEIF_HEADER];
	int magic_len;
	size_t size = 0, n = GD_HEIF_ALLOC_STEP;
	gdImagePtr im;
	int x, y;
	uint8_t *p, *row_start;
	int stride;

	magic_len = gdGetBuf(magic, GD_HEIF_HEADER, infile);
	if (magic_len != GD_HEIF_HEADER || !_gdHeifCheckBrand(magic, expected_brand)) {
		gd_error("gd-heif incorrect type of file\n");
		return NULL;
	}
	gdSeek(infile, 0);

	while (n == GD_HEIF_ALLOC_STEP) {
		temp = gdRealloc(filedata, size + GD_HEIF_ALLOC_STEP);
		if (temp) {
			filedata = temp;
			read = temp + size;
		} else {
			gdFree(filedata);
			gd_error("gd-heif decode realloc failed\n");
			return NULL;
		}

		n = gdGetBuf(read, GD_HEIF_ALLOC_STEP, infile);
		if (n > 0) {
			size += n;
		}
	}

	heif_ctx = heif_context_alloc();
	if (heif_ctx == NULL) {
		gd_error("gd-heif could not allocate context\n");
		gdFree(filedata);
		return NULL;
	}
	err = heif_context_read_from_memory_without_copy(heif_ctx, filedata, size, NULL);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif context creation failed\n");
		gdFree(filedata);
		heif_context_free(heif_ctx);
		return NULL;
	}

	heif_imhandle = NULL;
	err = heif_context_get_primary_image_handle(heif_ctx, &heif_imhandle);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif cannot retreive handle\n");
		gdFree(filedata);
		heif_context_free(heif_ctx);
		return NULL;
	}

	heif_im = NULL;
	heif_dec_opts = heif_decoding_options_alloc();
	if (heif_dec_opts == NULL) {
		gd_error("gd-heif could not allocate decode options\n");
		gdFree(filedata);
		heif_image_handle_release(heif_imhandle);
		heif_context_free(heif_ctx);
		return NULL;
	}

	heif_dec_opts->convert_hdr_to_8bit = GD_TRUE;
	heif_dec_opts->ignore_transformations = GD_TRUE;
	err = heif_decode_image(heif_imhandle, &heif_im, heif_colorspace_RGB, heif_chroma_interleaved_RGBA, heif_dec_opts);
	heif_decoding_options_free(heif_dec_opts);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif decoding failed\n");
		gdFree(filedata);
		heif_image_handle_release(heif_imhandle);
		heif_context_free(heif_ctx);
		return NULL;
	}

	width = heif_image_get_width(heif_im, heif_channel_interleaved);
	height = heif_image_get_height(heif_im, heif_channel_interleaved);

	im = gdImageCreateTrueColor(width, height);
	if (!im) {
		gdFree(filedata);
		heif_image_release(heif_im);
		heif_image_handle_release(heif_imhandle);
		heif_context_free(heif_ctx);
		return NULL;
	}
	rgba = (uint8_t *)heif_image_get_plane_readonly(heif_im, heif_channel_interleaved, &stride);
	if (!rgba) {
		gd_error("gd-heif cannot get image plane\n");
		gdFree(filedata);
		heif_image_release(heif_im);
		heif_image_handle_release(heif_imhandle);
		heif_context_free(heif_ctx);
		gdImageDestroy(im);
		return NULL;
	}
	row_start = rgba;
	for (y = 0, p = rgba; y < height; y++) {
		p = row_start;
		for (x = 0; x < width; x++) {
			uint8_t r = *(p++);
			uint8_t g = *(p++);
			uint8_t b = *(p++);
			uint8_t a = gdAlphaMax - (*(p++) >> 1);
			im->tpixels[y][x] = gdTrueColorAlpha(r, g, b, a);
		}
		row_start += stride;
	}
	gdFree(filedata);
	heif_image_release(heif_im);
	heif_image_handle_release(heif_imhandle);
	heif_context_free(heif_ctx);

	return im;
}

/*
  Function: gdImageCreateFromHeifCtx

    See <gdImageCreateFromHeif>.
*/
BGD_DECLARE(gdImagePtr) gdImageCreateFromHeifCtx(gdIOCtx *infile)
{
	return _gdImageCreateFromHeifCtx(infile, GD_HEIF_BRAND_AVIF | GD_HEIF_BRAND_MIF1 | GD_HEIF_BRAND_HEIC | GD_HEIF_BRAND_HEIX);
}


static struct heif_error _gdImageWriteHeif(struct heif_context *heif_ctx, const void *data, size_t size, void *userdata)
{
	ARG_NOT_USED(heif_ctx);
	gdIOCtx *outfile;
	struct heif_error err;

	outfile = (gdIOCtx *)userdata;

	gdPutBuf(data, size, outfile);

	err.code = heif_error_Ok;
	err.subcode = heif_suberror_Unspecified;
	err.message = "";

	return err;
}

/* returns GD_TRUE on success, GD_FALSE on failure */
static int _gdImageHeifCtx(gdImagePtr im, gdIOCtx *outfile, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	struct heif_context *heif_ctx;
	struct heif_encoder *heif_enc;
	struct heif_image *heif_im;
	struct heif_writer heif_wr;
	struct heif_error err;
	uint8_t *rgba;
	int x, y;
	uint8_t *p;
	uint8_t *row_start;
	int stride;
	if (im == NULL) {
		return GD_FALSE;
	}

	if (codec != GD_HEIF_CODEC_HEVC && codec != GD_HEIF_CODEC_AV1) {
		gd_error("Unsupported format by heif");
		return GD_FALSE;
	}

	if (!gdImageTrueColor(im)) {
		gd_error("Palette image not supported by heif\n");
		return GD_FALSE;
	}

	if (overflow2(gdImageSX(im), 4)) {
		return GD_FALSE;
	}

	if (overflow2(gdImageSX(im) * 4, gdImageSY(im))) {
		return GD_FALSE;
	}

	heif_ctx = heif_context_alloc();
	if (heif_ctx == NULL) {
		gd_error("gd-heif could not allocate context\n");
		return GD_FALSE;
	}
	err = heif_context_get_encoder_for_format(heif_ctx, (enum heif_compression_format)codec, &heif_enc);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif encoder acquisition failed (missing codec support?)\n");
		heif_context_free(heif_ctx);
		return GD_FALSE;
	}

	if (quality == 200) {
		err = heif_encoder_set_lossless(heif_enc, GD_TRUE);
	} else if (quality == -1) {
		err = heif_encoder_set_lossy_quality(heif_enc, 80);
	} else {
		err = heif_encoder_set_lossy_quality(heif_enc, quality);
	}
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif invalid quality number\n");
		heif_encoder_release(heif_enc);
		heif_context_free(heif_ctx);
		return GD_FALSE;
	}

	if (heif_get_version_number_major() >= 1 && heif_get_version_number_minor() >= 9) {
		err = heif_encoder_set_parameter_string(heif_enc, "chroma", chroma);
		if (err.code != heif_error_Ok) {
			gd_error("gd-heif invalid chroma subsampling parameter\n");
			heif_encoder_release(heif_enc);
			heif_context_free(heif_ctx);
			return GD_FALSE;
		}
	}

	err = heif_image_create(gdImageSX(im), gdImageSY(im), heif_colorspace_RGB, heif_chroma_interleaved_RGBA, &heif_im);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif image creation failed");
		heif_encoder_release(heif_enc);
		heif_context_free(heif_ctx);
		return GD_FALSE;
	}

	err = heif_image_add_plane(heif_im, heif_channel_interleaved, gdImageSX(im), gdImageSY(im), 32);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif cannot add image plane\n");
		heif_image_release(heif_im);
		heif_encoder_release(heif_enc);
		heif_context_free(heif_ctx);
		return GD_FALSE;
	}

	rgba = (uint8_t *)heif_image_get_plane_readonly(heif_im, heif_channel_interleaved, &stride);
	if (!rgba) {
		gd_error("gd-heif cannot get image plane\n");
		heif_image_release(heif_im);
		heif_encoder_release(heif_enc);
		heif_context_free(heif_ctx);
		return GD_FALSE;
	}
	row_start = rgba;
	for (y = 0; y < gdImageSY(im); y++) {
		p = row_start;
		for (x = 0; x < gdImageSX(im); x++) {
			int c;
			char a;
			c = im->tpixels[y][x];
			a = gdTrueColorGetAlpha(c);
			if (a == 127) {
				a = 0;
			} else {
				a = 255 - ((a << 1) + (a >> 6));
			}
			*(p++) = gdTrueColorGetRed(c);
			*(p++) = gdTrueColorGetGreen(c);
			*(p++) = gdTrueColorGetBlue(c);
			*(p++) = a;
		}
		row_start += stride;
	}
	err = heif_context_encode_image(heif_ctx, heif_im, heif_enc, NULL, NULL);
	heif_encoder_release(heif_enc);
	if (err.code != heif_error_Ok) {
		gd_error("gd-heif encoding failed\n");
		heif_image_release(heif_im);
		heif_context_free(heif_ctx);
		return GD_FALSE;
	}
	heif_wr.write = _gdImageWriteHeif;
	heif_wr.writer_api_version = 1;
	err = heif_context_write(heif_ctx, &heif_wr, (void *)outfile);

	heif_image_release(heif_im);
	heif_context_free(heif_ctx);

	return GD_TRUE;
}


/*
  Function: gdImageHeifCtx

    Write the image as HEIF data via a <gdIOCtx>. See <gdImageHeifEx>
    for more details.

  Parameters:

    im          - The image to write.
    outfile     - The output sink.
    quality     - Image quality.
    codec       - The output coding format.
    chroma      - The output chroma subsampling format.

  Returns:

    Nothing.
*/
BGD_DECLARE(void) gdImageHeifCtx(gdImagePtr im, gdIOCtx *outfile, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	_gdImageHeifCtx(im, outfile, quality, codec, chroma);
}

/*
  Function: gdImageHeifEx

    <gdImageHeifEx> outputs the specified image to the specified file in
    HEIF format. The file must be open for writing. Under MSDOS and
    all versions of Windows, it is important to use "wb" as opposed to
    simply "w" as the mode when opening the file, and under Unix there
    is no penalty for doing so. <gdImageHeifEx> does not close the file;
    your code must do so.

    If _quality_ is -1, a reasonable quality value (which should yield a
    good general quality / size tradeoff for most situations) is used. Otherwise
    _quality_ should be a value in the range 0-100, higher quality values
    usually implying both higher quality and larger image sizes or 200, for
    lossless codec.

  Variants:

    <gdImageHeifCtx> stores the image using a <gdIOCtx> struct.

    <gdImageHeifPtrEx> stores the image to RAM.

  Parameters:

    im          - The image to save.
    outFile     - The FILE pointer to write to.
    quality     - Codec quality (0-100).
    codec       - The output coding format.
    chroma      - The output chroma subsampling format.

  Returns:

    Nothing.
*/
BGD_DECLARE(void) gdImageHeifEx(gdImagePtr im, FILE *outFile, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	gdIOCtx *out = gdNewFileCtx(outFile);
	if (out == NULL) {
		return;
	}
	_gdImageHeifCtx(im, out, quality, codec, chroma);
	out->gd_free(out);
}

/*
  Function: gdImageHeif

    Variant of <gdImageHeifEx> which uses the default quality (-1), the
    default codec (GD_HEIF_Codec_HEVC) and the default chroma
    subsampling (GD_HEIF_CHROMA_444).

  Parameters:

    im      - The image to save
    outFile - The FILE pointer to write to.

  Returns:

    Nothing.
*/
BGD_DECLARE(void) gdImageHeif(gdImagePtr im, FILE *outFile)
{
	gdIOCtx *out = gdNewFileCtx(outFile);
	if (out == NULL) {
		return;
	}
	_gdImageHeifCtx(im, out, -1, GD_HEIF_CODEC_HEVC, GD_HEIF_CHROMA_444);
	out->gd_free(out);
}

/*
  Function: gdImageHeifPtr

    See <gdImageHeifEx>.
*/
BGD_DECLARE(void *) gdImageHeifPtr(gdImagePtr im, int *size)
{
	void *rv;
	gdIOCtx *out = gdNewDynamicCtx(2048, NULL);
	if (out == NULL) {
		return NULL;
	}
	if (_gdImageHeifCtx(im, out, -1, GD_HEIF_CODEC_HEVC, GD_HEIF_CHROMA_444)) {
		rv = gdDPExtractData(out, size);
	} else {
		rv = NULL;
	}
	out->gd_free(out);

	return rv;
}

/*
  Function: gdImageHeifPtrEx

    See <gdImageHeifEx>.
*/
BGD_DECLARE(void *) gdImageHeifPtrEx(gdImagePtr im, int *size, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	void *rv;
	gdIOCtx *out = gdNewDynamicCtx(2048, NULL);
	if (out == NULL) {
		return NULL;
	}
	if (_gdImageHeifCtx(im, out, quality, codec, chroma)) {
		rv = gdDPExtractData(out, size);
	} else {
		rv = NULL;
	}
	out->gd_free(out);
	return rv;
}

#else /* HAVE_LIBHEIF */

static void _noHeifError(void)
{
	gd_error("HEIF image support has been disabled\n");
}

BGD_DECLARE(gdImagePtr) gdImageCreateFromHeif(FILE *inFile)
{
	_noHeifError();
	return NULL;
}

BGD_DECLARE(gdImagePtr) gdImageCreateFromHeifPtr(int size, void *data)
{
	_noHeifError();
	return NULL;
}

BGD_DECLARE(gdImagePtr) gdImageCreateFromHeifCtx(gdIOCtx *infile)
{
	_noHeifError();
	return NULL;
}

BGD_DECLARE(void) gdImageHeifCtx(gdImagePtr im, gdIOCtx *outfile, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	_noHeifError();
}

BGD_DECLARE(void) gdImageHeifEx(gdImagePtr im, FILE *outFile, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	_noHeifError();
}

BGD_DECLARE(void) gdImageHeif(gdImagePtr im, FILE *outFile)
{
	_noHeifError();
}

BGD_DECLARE(void *) gdImageHeifPtr(gdImagePtr im, int *size)
{
	_noHeifError();
	return NULL;
}

BGD_DECLARE(void *) gdImageHeifPtrEx(gdImagePtr im, int *size, int quality, gdHeifCodec codec, gdHeifChroma chroma)
{
	_noHeifError();
	return NULL;
}

#endif /* HAVE_LIBHEIF */
