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

static void
register_custom_tiff_tags(TIFF *tif)
{
    static const TIFFFieldInfo custom_fields[] = {
        {TIFFTAG_FORWARDMATRIX1, -1, -1, TIFF_SRATIONAL, FIELD_CUSTOM, 1, 1,
            "ForwardMatrix1"},
    };

    // Add missing dng fields
    TIFFMergeFieldInfo(tif, custom_fields,
        sizeof(custom_fields) / sizeof(custom_fields[0]));
}

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;

    TIFF *im = TIFFOpen(filename, "r");

    if (im == NULL) {
        fprintf(stderr, "Could not open tiff");
        exit(1);
    }

    if (TIFFGetField(im, TIFFTAG_MAKE, &temp) != 1) {
        err("failed to read TIFFTAG_MAKE");
    }
    imagedata.make = strdup(temp);

    if (TIFFGetField(im, TIFFTAG_MODEL, &temp) != 1) {
        err("failed to read TIFFTAG_MODEL");
    }
    imagedata.model = strdup(temp);

    if (TIFFGetField(im, TIFFTAG_SOFTWARE, &temp) != 1) {
        err("failed to read TIFFTAG_SOFTWARE");
    }
    imagedata.software = strdup(temp);

    if (TIFFGetField(im, TIFFTAG_DATETIME, &temp) != 1) {
        err("failed to read TIFFTAG_DATETIME");
    }
    imagedata.datetime = strdup(temp);

    if (TIFFGetField(im, TIFFTAG_ORIENTATION, &imagedata.orientation) != 1) {
        err("failed to read TIFFTAG_ORIENTATION");
    }

    // Get the EXIF directory for the rest of the metadata
    if (TIFFGetField(im, TIFFTAG_EXIFIFD, &exif_offset) != 1) {
        err("failed to read TIFFTAG_EXIFIFD");
    }

    // Read the actual picture instead of the embedded thumbnail
    if (TIFFGetField(im, TIFFTAG_SUBIFD, &subifd_count, &subifd_offsets)) {
        if (subifd_count > 0) {
            TIFFSetSubDirectory(im, subifd_offsets[0]);
        }
    }

    if (TIFFGetField(im, TIFFTAG_IMAGELENGTH, &imagedata.height) != 1) {
        err("failed to read TIFFTAG_IMAGELENGTH");
    }
    if (TIFFGetField(im, TIFFTAG_IMAGEWIDTH, &imagedata.width) != 1) {
        err("failed to read TIFFTAG_IMAGEWIDTH");
    }
    if (TIFFGetField(im, TIFFTAG_BITSPERSAMPLE, &imagedata.bitspersample) != 1) {
        err("failed to read TIFFTAG_BITSPERSAMPLE");
    }

    // Read the exif directory
    TIFFReadEXIFDirectory(im, exif_offset);

    if (TIFFGetField(im, EXIFTAG_EXPOSUREPROGRAM, &imagedata.exposure_program) != 1) {
        err("failed to read EXIFTAG_EXPOSUREPROGRAM");
    }

    if (TIFFGetField(im, EXIFTAG_EXPOSURETIME, &imagedata.exposure_time) != 1) {
        err("failed to read EXIFTAG_EXPOSURETIME");
    }

    if (TIFFGetField(im, EXIFTAG_PHOTOGRAPHICSENSITIVITY, &value_count, &short_array) != 1) {
        err("failed to read EXIFTAG_PHOTOGRAPHICSENSITIVITY");
    }
    imagedata.isospeed = short_array[0];

    if (TIFFGetField(im, EXIFTAG_FNUMBER, &imagedata.fnumber) != 1) {
        err("failed to read EXIFTAG_FNUMBER");
    }

    if (TIFFGetField(im, EXIFTAG_FOCALLENGTH, &imagedata.focal_length) != 1) {
        err("failed to read EXIFTAG_FOCALLENGTH");
    }

    if (TIFFGetField(im, EXIFTAG_FOCALLENGTHIN35MMFILM, &imagedata.focal_length_35mm) != 1) {
        err("failed to read EXIFTAG_FOCALLENGTHIN35MMFILM");
    }

    if (TIFFGetField(im, EXIFTAG_FLASH, &imagedata.flash) != 1) {
        err("failed to read EXIFTAG_FLASH");
    }
    TIFFClose(im);
    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()
{
    TIFFSetTagExtender(register_custom_tiff_tags);
}

void
postprocess_internal(char *burst_dir, char *target_dir)
{
    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];

    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);
    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);
        snprintf(outpath, sizeof(outpath), "%s.%d.jpg", target_dir, burst_size);
        if (first) {
            first = 0;
            imagedata = read_exif(path);
            exif = create_exif(imagedata);
        }
        decoded = debayer_file(path);
        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]);
    }
    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");
    snprintf(stackpath, sizeof(stackpath), "%s.stacked.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;
    stacker_t *stacker = stacker_create(verbose);

    // Parse exif data from original file
    timer = clock();
    imagedata = read_exif(in_path);
    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");
    }
}