//
// benchmark-raw.c: measure loading times of raw images and their thumbnails
//
// This is a tool to help decide on criteria for fast thumbnail extraction.
//
// Copyright (c) 2023, Přemysl Eric Janouch <p@janouch.name>
//
// Permission to use, copy, modify, and/or distribute this software for any
// purpose with or without fee is hereby granted.
//
// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
// WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
// MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY
// SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
// WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION
// OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF OR IN
// CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
//

#include <gio/gio.h>
#include <jv.h>
#include <libraw.h>

#include <stdbool.h>
#include <time.h>

#if LIBRAW_VERSION < LIBRAW_MAKE_VERSION(0, 21, 0)
#error LibRaw 0.21.0 or newer is required.
#endif

#include "fiv-io.h"
#include "fiv-thumbnail.h"

// --- Analysis ----------------------------------------------------------------
// Functions duplicated from info.h and benchmark-io.c.

static jv
add_to_subarray(jv o, const char *key, jv value)
{
	// Invalid values are not allocated, and we use up any valid one.
	// Beware that jv_get() returns jv_null() rather than jv_invalid().
	// Also, the header comment is lying, jv_is_valid() doesn't unreference.
	jv a = jv_object_get(jv_copy(o), jv_string(key));
	return jv_set(o, jv_string(key),
		jv_is_valid(a) ? jv_array_append(a, value) : JV_ARRAY(value));
}

static jv
add_warning(jv o, const char *message)
{
	return add_to_subarray(o, "warnings", jv_string(message));
}

static jv
add_error(jv o, const char *message)
{
	return jv_object_set(o, jv_string("error"), jv_string(message));
}

static double
timestamp(void)
{
	struct timespec ts;
	clock_gettime(CLOCK_MONOTONIC, &ts);
	return ts.tv_sec + ts.tv_nsec / 1.e9;
}

// --- Raw image files ---------------------------------------------------------

static bool extract_mode = false;

// Copied function from fiv-thumbnail.c.
static FivIoImage *
orient_thumbnail(FivIoImage *image)
{
	if (image->orientation <= FivIoOrientation0)
		return image;

	double w = 0, h = 0;
	cairo_matrix_t matrix =
		fiv_io_orientation_apply(image, image->orientation, &w, &h);
	FivIoImage *oriented = fiv_io_image_new(image->format, w, h);
	if (!oriented) {
		g_warning("image allocation failure");
		return image;
	}

	cairo_surface_t *surface = fiv_io_image_to_surface_noref(oriented);
	cairo_t *cr = cairo_create(surface);
	cairo_surface_destroy(surface);

	surface = fiv_io_image_to_surface(image);
	cairo_set_source_surface(cr, surface, 0, 0);
	cairo_surface_destroy(surface);
	cairo_pattern_set_matrix(cairo_get_source(cr), &matrix);
	cairo_set_operator(cr, CAIRO_OPERATOR_SOURCE);
	cairo_paint(cr);
	cairo_destroy(cr);
	return oriented;
}

// Modified function from fiv-thumbnail.c.
static FivIoImage *
adjust_thumbnail(FivIoImage *thumbnail, double row_height)
{
	// Hardcode orientation.
	FivIoOrientation orientation = thumbnail->orientation;

	double w = 0, h = 0;
	cairo_matrix_t matrix =
		fiv_io_orientation_apply(thumbnail, orientation, &w, &h);

	double scale_x = 1;
	double scale_y = 1;
	if (w > FIV_THUMBNAIL_WIDE_COEFFICIENT * h) {
		scale_x = FIV_THUMBNAIL_WIDE_COEFFICIENT * row_height / w;
		scale_y = round(scale_x * h) / h;
	} else {
		scale_y = row_height / h;
		scale_x = round(scale_y * w) / w;
	}

	// NOTE: Ignoring renderable images.

	if (orientation <= FivIoOrientation0 && scale_x == 1 && scale_y == 1)
		return fiv_io_image_ref(thumbnail);

	cairo_format_t format = thumbnail->format;
	int projected_width = round(scale_x * w);
	int projected_height = round(scale_y * h);
	FivIoImage *scaled = fiv_io_image_new(
		(format == CAIRO_FORMAT_RGB24 || format == CAIRO_FORMAT_RGB30)
			? CAIRO_FORMAT_RGB24
			: CAIRO_FORMAT_ARGB32,
		projected_width, projected_height);
	if (!scaled)
		return fiv_io_image_ref(thumbnail);

	cairo_surface_t *surface = fiv_io_image_to_surface_noref(scaled);
	cairo_t *cr = cairo_create(surface);
	cairo_surface_destroy(surface);

	cairo_scale(cr, scale_x, scale_y);

	surface = fiv_io_image_to_surface_noref(thumbnail);
	cairo_set_source_surface(cr, surface, 0, 0);
	cairo_surface_destroy(surface);

	cairo_pattern_t *pattern = cairo_get_source(cr);
	// CAIRO_FILTER_BEST, for some reason, works bad with CAIRO_FORMAT_RGB30.
	cairo_pattern_set_filter(pattern, CAIRO_FILTER_GOOD);
	cairo_pattern_set_extend(pattern, CAIRO_EXTEND_PAD);
	cairo_pattern_set_matrix(pattern, &matrix);

	cairo_set_operator(cr, CAIRO_OPERATOR_SOURCE);
	cairo_paint(cr);

	// NOTE: Ignoring silent cairo errors.

	cairo_destroy(cr);
	return scaled;
}

// Copied function from fiv-thumbnail.c.
// LibRaw does a weird permutation here, so follow the documentation,
// which assumes that mirrored orientations never happen.
static FivIoOrientation
extract_libraw_unflip(int flip)
{
	switch (flip) {
	break; case 0:
		return FivIoOrientation0;
	break; case 3:
		return FivIoOrientation180;
	break; case 5:
		return FivIoOrientation270;
	break; case 6:
		return FivIoOrientation90;
	break; default:
		return FivIoOrientationUnknown;
	}
}

// Modified function from fiv-thumbnail.c.
static FivIoImage *
extract_libraw_bitmap(libraw_processed_image_t *image, int flip)
{
	// Anything else is extremely rare.
	if (image->colors != 3 || image->bits != 8)
		return NULL;

	FivIoImage *I = fiv_io_image_new(
		CAIRO_FORMAT_RGB24, image->width, image->height);
	if (!I)
		return NULL;

	guint32 *out = (guint32 *) I->data;
	const unsigned char *in = image->data;
	for (guint64 i = 0; i < image->width * image->height; in += 3)
		out[i++] = in[0] << 16 | in[1] << 8 | in[2];

	I->orientation = extract_libraw_unflip(flip);
	return I;
}

static jv
process_thumbnail(
	jv o, FivIoOpenContext *ctx, libraw_data_t *iprc, int i)
{
	double since = timestamp();

	int err = 0;
	if ((err = libraw_unpack_thumb_ex(iprc, i))) {
		if (err != LIBRAW_NO_THUMBNAIL)
			o = add_warning(o, libraw_strerror(err));
		return o;
	}

	libraw_thumbnail_item_t *item = iprc->thumbs_list.thumblist + i;
	jv to = JV_OBJECT(
		jv_string("width"), jv_number(item->twidth),
		jv_string("height"), jv_number(item->theight));

	libraw_processed_image_t *image = libraw_dcraw_make_mem_thumb(iprc, &err);
	if (!image) {
		o = add_warning(o, libraw_strerror(err));
		goto fail;
	}

	FivIoImage *I = NULL;
	FivIoOrientation orientation = 0;
	switch (image->type) {
	break; case LIBRAW_IMAGE_JPEG:
		I = fiv_io_open_from_data(
			(const char *) image->data, image->data_size, ctx, NULL);
		orientation = I->orientation;
	break; case LIBRAW_IMAGE_BITMAP:
		I = extract_libraw_bitmap(image, item->tflip);
		orientation = I->orientation;
	break; default:
		o = add_warning(o, "unsupported embedded thumbnail");
	}
	if (!I)
		goto fail_render;

	if (item->tflip != 0xffff &&
		extract_libraw_unflip(item->tflip) != orientation) {
		gchar *m = g_strdup_printf("Orientation mismatch: tflip %d, Exif %d",
			extract_libraw_unflip(item->tflip), orientation);
		o = add_warning(o, m);
		g_free(m);
	}

	double width = 0, height = 0;
	fiv_io_orientation_dimensions(I, orientation, &width, &height);
	to = jv_set(to, jv_string("width"), jv_number(width));
	to = jv_set(to, jv_string("height"), jv_number(height));

	to = jv_set(to, jv_string("pixels_percent"),
		jv_number(100 * (width * height) /
			((float) iprc->sizes.iwidth * iprc->sizes.iheight)));

	float main_ratio = (float) iprc->sizes.iwidth / iprc->sizes.iheight;
	to = jv_set(to, jv_string("ratio_difference_percent"),
		jv_number(fabs((main_ratio - width / height) * 100)));

	// Resize, hardcode orientation. This may take a long time.
	to = jv_set(to, jv_string("duration_decode_ms"),
		jv_number((timestamp() - since) * 1000));
	fiv_io_image_unref(adjust_thumbnail(I, 256.));
	to = jv_set(to, jv_string("duration_ms"),
		jv_number((timestamp() - since) * 1000));

	// Luckily, large thumbnails are typically JPEGs, which don't need encoding.
	gchar *path = NULL;
	GError *error = NULL;
	if (extract_mode && (path = g_filename_from_uri(ctx->uri, NULL, &error))) {
		gchar *thumbnail_path = NULL;
		if (image->type == LIBRAW_IMAGE_JPEG) {
			thumbnail_path = g_strdup_printf("%s.thumb.%d.jpg", path, i);
			g_file_set_contents(thumbnail_path,
				(const char *) image->data, image->data_size, &error);
		} else {
			thumbnail_path = g_strdup_printf("%s.thumb.%d.webp", path, i);
			I = orient_thumbnail(I);
			fiv_io_save(I, I, NULL, thumbnail_path, &error);
		}

		g_clear_pointer(&thumbnail_path, g_free);
		g_clear_pointer(&path, g_free);
	}
	if (error) {
		o = add_warning(o, error->message);
		g_clear_error(&error);
	}

	g_clear_pointer(&I, fiv_io_image_unref);
fail_render:
	libraw_dcraw_clear_mem(image);
fail:
	return add_to_subarray(o, "thumbnails", to);
}

static jv
process_raw(jv o, const char *filename, const uint8_t *data, size_t len)
{
	libraw_data_t *iprc = libraw_init(LIBRAW_OPIONS_NO_DATAERR_CALLBACK);
	if (!iprc)
		return add_error(o, "failed to obtain a LibRaw handle");

	// First, bail out if this isn't a raw image file.
	int err = 0;
	if ((err = libraw_open_buffer(iprc, data, len)) ||
		(err = libraw_adjust_sizes_info_only(iprc))) {
		o = add_error(o, libraw_strerror(err));
		goto fail;
	}

	// Run our entire stack, like the render() function in fiv-thumbnail.c does.
	// Note that this may use the TIFF/EP shortcut code.
	double since = timestamp();
	GFile *file = g_file_new_for_commandline_arg(filename);
	FivIoOpenContext ctx = {
		.uri = g_file_get_uri(file),
		.screen_profile = fiv_io_profile_new_sRGB(),
		.screen_dpi = 96,
		.warnings = g_ptr_array_new_with_free_func(g_free),
	};
	g_clear_object(&file);

	// This is really slow, let's decouple the mode from measurement a bit.
	if (!extract_mode) {
		GError *error = NULL;
		FivIoImage *image =
			fiv_io_open_from_data((const char *) data, len, &ctx, &error);
		if (!image) {
			o = add_error(o, error->message);
			g_error_free(error);
			goto fail_context;
		}

		// Resize, hardcode orientation. This may take a long time.
		o = jv_set(o, jv_string("duration_decode_ms"),
			jv_number((timestamp() - since) * 1000));
		fiv_io_image_unref(adjust_thumbnail(image, 256.));
		g_clear_pointer(&image, fiv_io_image_unref);

		o = jv_set(o, jv_string("duration_ms"),
			jv_number((timestamp() - since) * 1000));
	}

	o = jv_set(o, jv_string("thumbnails"), jv_array());
	for (int i = 0; i < iprc->thumbs_list.thumbcount; i++)
		o = process_thumbnail(o, &ctx, iprc, i);

fail_context:
	g_free((char *) ctx.uri);
	if (ctx.screen_profile)
		fiv_io_profile_free(ctx.screen_profile);

	for (guint i = 0; i < ctx.warnings->len; i++)
		o = add_warning(o, ctx.warnings->pdata[i]);
	g_ptr_array_free(ctx.warnings, TRUE);

fail:
	libraw_close(iprc);
	return o;
}

// --- I/O ---------------------------------------------------------------------

static jv
do_file(const char *filename, jv o)
{
	const char *err = NULL;
	FILE *fp = fopen(filename, "rb");
	if (!fp) {
		err = strerror(errno);
		goto error;
	}

	uint8_t *data = NULL, buf[256 << 10];
	size_t n, len = 0;
	while ((n = fread(buf, sizeof *buf, sizeof buf / sizeof *buf, fp))) {
		data = realloc(data, len + n);
		memcpy(data + len, buf, n);
		len += n;
	}
	if (ferror(fp)) {
		err = strerror(errno);
		goto error_read;
	}

	o = process_raw(o, filename, data, len);

error_read:
	fclose(fp);
	free(data);
error:
	if (err)
		o = add_error(o, err);
	return o;
}

int
main(int argc, char *argv[])
{
	// We don't need to call gdk_cairo_surface_create_from_pixbuf() here,
	// so don't bother initializing GDK.

	fiv_io_profile_init();

	// A mode to just extract all thumbnails to files for closer inspection.
	extract_mode = !!getenv("BENCHMARK_RAW_EXTRACT");

	// XXX: Can't use `xargs -P0`, there's a risk of non-atomic writes.
	// Usage: find . -type f -print0 | xargs -0 ./benchmark-raw
	for (int i = 1; i < argc; i++) {
		const char *filename = argv[i];

		jv o = jv_object();
		o = jv_object_set(o, jv_string("filename"), jv_string(filename));
		o = do_file(filename, o);
		jv_dumpf(o, stdout, 0 /* Might consider JV_PRINT_SORTED. */);
		fputc('\n', stdout);
	}
	return 0;
}