#include <assert.h>
#include <string.h>
#include <dirent.h>
#include <libraw/libraw.h>
#include <jpeglib.h>
#include <libexif/exif-data.h>
#include <time.h>
#include <unistd.h>
#include <libdng.h>
#include "postprocess.h"
#include "stacker.h"

libraw_processed_image_t *
debayer_file(char *filename)
{
	libraw_data_t *raw;
	int ret;

	raw = libraw_init(0);
	if (libraw_open_file(raw, filename) != LIBRAW_SUCCESS) {
		err("could not open file");
		libraw_close(raw);
	}

	if (libraw_unpack(raw) != LIBRAW_SUCCESS) {
		err("could not unpack file");
		libraw_close(raw);
	}

	raw->params.no_auto_bright = 1;

	if (libraw_dcraw_process(raw) != LIBRAW_SUCCESS) {
		err("could not process file");
		libraw_free_image(raw);
		libraw_close(raw);
	}

	libraw_processed_image_t *proc_img = libraw_dcraw_make_mem_image(raw, &ret);
	libraw_free_image(raw);

	if (!proc_img) {
		err("could not export image");
		libraw_close(raw);
	}
	libraw_recycle(raw);
	libraw_close(raw);
	return proc_img;
}

void
save_jpeg(char *path, libraw_processed_image_t *data, int quality, ExifData *exif)
{
	unsigned char *exif_data;
	unsigned int exif_len;
	FILE *out = fopen(path, "wb");
	if (!out) {
		err("could not open target file");
		exit(1);
	}

	struct jpeg_compress_struct jpeg;
	struct jpeg_error_mgr error_mgr;
	jpeg.err = jpeg_std_error(&error_mgr);
	jpeg_create_compress(&jpeg);
	jpeg_stdio_dest(&jpeg, out);
	jpeg.image_width = data->width;
	jpeg.image_height = data->height;
	jpeg.input_components = 3;
	jpeg.in_color_space = JCS_RGB;
	jpeg_set_defaults(&jpeg);
	jpeg_set_quality(&jpeg, quality, 1);
	jpeg_start_compress(&jpeg, 1);

	// Write exif
	exif_data_save_data(exif, &exif_data, &exif_len);
	jpeg_write_marker(&jpeg, JPEG_APP1, exif_data, exif_len);

	// Write image data
	JSAMPROW row_pointer;
	int row_stride = jpeg.image_width * jpeg.input_components;

	while (jpeg.next_scanline < jpeg.image_height) {
		row_pointer = (JSAMPROW) &data->data[jpeg.next_scanline * row_stride];
		jpeg_write_scanlines(&jpeg, &row_pointer, 1);
	}

	jpeg_finish_compress(&jpeg);
	jpeg_destroy_compress(&jpeg);
	fclose(out);
}

static ExifEntry *
init_tag(ExifData *exif, ExifIfd ifd, ExifTag tag)
{
	ExifEntry *entry;
	/* Return an existing tag if one exists */
	if (!((entry = exif_content_get_entry(exif->ifd[ifd], tag)))) {
		/* Allocate a new entry */
		entry = exif_entry_new();
		assert(entry != NULL); /* catch an out of memory condition */
		entry->tag = tag; /* tag must be set before calling
				 exif_content_add_entry */

		/* Attach the ExifEntry to an IFD */
		exif_content_add_entry(exif->ifd[ifd], entry);

		/* Allocate memory for the entry and fill with default data */
		exif_entry_initialize(entry, tag);

		/* Ownership of the ExifEntry has now been passed to the IFD.
		 * One must be very careful in accessing a structure after
		 * unref'ing it; in this case, we know "entry" won't be freed
		 * because the reference count was bumped when it was added to
		 * the IFD.
		 */
		exif_entry_unref(entry);
	}
	return entry;
}

void
exif_set_string(ExifEntry *ed, const char *s, size_t size)
{
	if (ed->data) {
		free(ed->data);
	}
	ed->components = size + 1;
	ed->size = sizeof(char) * ed->components;
	ed->data = (unsigned char *) malloc(ed->size);
	if (!ed->data) {
		err("Could not allocate exif string");
		exit(1);
	}
	strncpy((char *) ed->data, (char *) s, size);
	exif_entry_fix(ed);
}

ExifData *
create_exif(struct Imagedata data)
{
	ExifEntry *entry;
	ExifRational rational;
	long denominator = 100000;
	ExifData *exif = exif_data_new();
	if (!exif) {
		err("Could not initialize libexif");
	}
	exif_data_set_option(exif, EXIF_DATA_OPTION_FOLLOW_SPECIFICATION);
	exif_data_set_data_type(exif, EXIF_DATA_TYPE_COMPRESSED);
	exif_data_set_byte_order(exif, EXIF_BYTE_ORDER_INTEL);
	exif_data_fix(exif);

	// Width
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_PIXEL_X_DIMENSION);
	exif_set_long(entry->data, EXIF_BYTE_ORDER_INTEL, data.width);

	// Height
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_PIXEL_Y_DIMENSION);
	exif_set_long(entry->data, EXIF_BYTE_ORDER_INTEL, data.height);

	// Colorspace, 1=sRGB
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_COLOR_SPACE);
	exif_set_long(entry->data, EXIF_BYTE_ORDER_INTEL, 1);

	// Exposure program, enum
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_EXPOSURE_PROGRAM);
	exif_set_long(entry->data, EXIF_BYTE_ORDER_INTEL, data.exposure_program);

	// Camera make
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_MAKE);
	exif_set_string(entry, data.make, strlen(data.make));

	// Camera model
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_MODEL);
	exif_set_string(entry, data.model, strlen(data.model));

	// Processing software
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_SOFTWARE);
	exif_set_string(entry, data.software, strlen(data.software));

	// Various datetime fields
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_DATE_TIME);
	exif_set_string(entry, data.datetime, strlen(data.datetime));
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_DATE_TIME_DIGITIZED);
	exif_set_string(entry, data.datetime, strlen(data.datetime));
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_DATE_TIME_ORIGINAL);
	exif_set_string(entry, data.datetime, strlen(data.datetime));

	// Exposure time
	rational.numerator = (long) ((double) data.exposure_time * denominator);
	rational.denominator = denominator;
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_EXPOSURE_TIME);
	exif_set_rational(entry->data, EXIF_BYTE_ORDER_INTEL, rational);

	// fnumber
	rational.numerator = (long) ((double) data.fnumber * denominator);
	rational.denominator = denominator;
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_FNUMBER);
	exif_set_rational(entry->data, EXIF_BYTE_ORDER_INTEL, rational);

	// focal length
	rational.numerator = (long) ((double) data.focal_length * denominator);
	rational.denominator = denominator;
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_FOCAL_LENGTH);
	exif_set_rational(entry->data, EXIF_BYTE_ORDER_INTEL, rational);

	// focal length, 35mm equiv
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_FOCAL_LENGTH_IN_35MM_FILM);
	exif_set_short(entry->data, EXIF_BYTE_ORDER_INTEL, data.focal_length_35mm);

	// ISO
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_ISO_SPEED_RATINGS);
	exif_set_long(entry->data, EXIF_BYTE_ORDER_INTEL, data.isospeed);

	// Flash
	entry = init_tag(exif, EXIF_IFD_EXIF, EXIF_TAG_FLASH);
	exif_set_short(entry->data, EXIF_BYTE_ORDER_INTEL, data.flash);
	return exif;
}

struct Imagedata
read_exif(char *filename)
{
	struct Imagedata imagedata;
	uint16_t subifd_count;
	uint32_t *subifd_offsets;
	uint32_t exif_offset;
	uint32_t value_count;
	uint16_t *short_array;
	char *temp;

	libdng_info dng = {0};
	if (libdng_read(&dng, filename) != 1) {
		err("Could not load DNG file");
	}
	imagedata.make = strdup(dng.camera_make);
	imagedata.model = strdup(dng.camera_model);
	imagedata.datetime = malloc(20);
	imagedata.software = strdup(dng.software);
	strftime(imagedata.datetime, 20, "%Y:%m:%d %H:%M:%S", &dng.datetime);
	imagedata.orientation = dng.orientation;
	imagedata.width = dng.width;
	imagedata.height = dng.height;
	imagedata.bitspersample = dng.bit_depth;
	imagedata.exposure_time = dng.exposure_time;
	imagedata.exposure_program = dng.exposure_program;
	imagedata.fnumber = dng.fnumber;
	imagedata.isospeed = (uint16_t) dng.iso;
	imagedata.focal_length = dng.focal_length;
	imagedata.focal_length_35mm = (uint16_t) (dng.focal_length * dng.crop_factor);
	imagedata.dist_a = dng.distortion_a;
	imagedata.dist_b = dng.distortion_b;
	imagedata.dist_c = dng.distortion_c;
	imagedata.vignette_k1 = dng.vignette_k1;
	imagedata.vignette_k2 = dng.vignette_k2;
	imagedata.vignette_k3 = dng.vignette_k3;
	imagedata.flash = 0;
	return imagedata;
}

static int
filter_dng(const struct dirent *dir)
{
	if (!dir)
		return 0;

	if (dir->d_type == DT_REG) {
		const char *ext = strrchr(dir->d_name, '.');
		if ((!ext) || (ext == dir->d_name))
			return 0;
		else {
			if (strcmp(ext, ".dng") == 0)
				return 1;
		}
	}

	return 0;
}

void
postprocess_setup(void)
{
	libdng_init();
}

void
postprocess_internal(char *burst_dir, char *target_dir, int keep)
{
	struct dirent **namelist;
	int burst_size;
	int first = 1;
	struct Imagedata imagedata;
	libraw_processed_image_t *decoded;
	ExifData *exif;
	char path[512];
	char outpath[512];
	char stackpath[512];

	nice(19);

	burst_size = scandir(burst_dir, &namelist, filter_dng, alphasort);
	if (burst_size < 0) {
		printf("oop\n");
		exit(1);
	}
	if (burst_size == 0) {
		printf("No DNG files found\n");
		exit(1);
	}
	stacker_t *stacker = stacker_create(1, imagedata);
	while (burst_size--) {
		fprintf(stderr, "DEBAYER %s\n", namelist[burst_size]->d_name);
		snprintf(path, sizeof(path), "%s/%s", burst_dir, namelist[burst_size]->d_name);
		if (keep)
			snprintf(outpath, sizeof(outpath), "%s.%d.jpg", target_dir, burst_size);
		else
			snprintf(outpath, sizeof(outpath), "%s.jpg", target_dir);

		if (first) {
			imagedata = read_exif(path);
			exif = create_exif(imagedata);
		}
		decoded = debayer_file(path);
		if (keep || first) {
			fprintf(stderr, "JPEG    %s\n", namelist[burst_size]->d_name);
			save_jpeg(outpath, decoded, 90, exif);
		}
		fprintf(stderr, "CONVERG %s\n", namelist[burst_size]->d_name);
		stacker_add_image(stacker, decoded->data, decoded->width, decoded->height);
		free(namelist[burst_size]);

		if (first) {
			first = 0;
		}
	}
	free(namelist);

	fprintf(stderr, "STACK   stacked.jpg\n");
	// Convert opencv result to a libraw struct and feed it in the jpeg encoder with the original exif data
	char *stacked = stacker_get_result(stacker);
	decoded->width = stacker_get_width(stacker);
	decoded->height = stacker_get_height(stacker);
	int result_size = decoded->width * decoded->height * 3;
	libraw_processed_image_t *stacked_data = (libraw_processed_image_t *) malloc(
		sizeof(libraw_processed_image_t) + result_size);
	memset(stacked_data, 0, sizeof(libraw_processed_image_t));
	stacked_data->colors = 3;
	stacked_data->width = decoded->width;
	stacked_data->height = decoded->height;
	memmove(stacked_data->data, stacked, result_size);

	// Save the final file
	fprintf(stderr, "JPEG    stacked.jpg\n");
	if (keep)
		snprintf(stackpath, sizeof(stackpath), "%s.stacked.jpg", target_dir);
	else
		snprintf(stackpath, sizeof(stackpath), "%s.jpg", target_dir);

	save_jpeg(stackpath, stacked_data, 90, exif);
}

void
postprocess_single(char *in_path, char *out_path, int quality, int verbose)
{
	libraw_processed_image_t *decoded;
	libraw_processed_image_t *processed_data;
	struct Imagedata imagedata;
	ExifData *exif;
	int width, height, result_size;
	char *result;
	clock_t timer;
	imagedata = read_exif(in_path);
	stacker_t *stacker = stacker_create(verbose, imagedata);

	// Give the operating system more cpu time
	nice(19);

	// Parse exif data from original file
	timer = clock();
	exif = create_exif(imagedata);
	if (verbose) {
		printf("[%.1fms] %s\n", (float) (clock() - timer) / CLOCKS_PER_SEC * 1000, "exif read");
	}

	// Convert the sensor data to rgb
	timer = clock();
	decoded = debayer_file(in_path);
	if (verbose) {
		printf("[%.1fms] %s\n", (float) (clock() - timer) / CLOCKS_PER_SEC * 1000, "debayer");
	}

	// Run the opencv postprocessing on the single frame
	result = stacker_postprocess(stacker, decoded->data, decoded->width, decoded->height);
	width = stacker_get_width(stacker);
	height = stacker_get_height(stacker);

	// Export the final jpeg with the original exif data
	timer = clock();
	result_size = width * height * 3;
	processed_data = (libraw_processed_image_t *) malloc(sizeof(libraw_processed_image_t) + result_size);
	memset(processed_data, 0, sizeof(libraw_processed_image_t));
	processed_data->colors = 3;
	processed_data->width = decoded->width;
	processed_data->height = decoded->height;
	memmove(processed_data->data, result, result_size);
	save_jpeg(out_path, processed_data, quality, exif);
	if (verbose) {
		printf("[%.1fms] %s\n", (float) (clock() - timer) / CLOCKS_PER_SEC * 1000, "export");
	}
}