#include <stdio.h>

// for error recovery
#include <setjmp.h>

#include <iostream>
#include <fstream>
#include <sstream>
#include <vector>

#include "jpeglib.h"

#define STB_IMAGE_IMPLEMENTATION
#include "stb_image.h"

#include "image.h"
#include "exif.h"

using namespace std;
using namespace haader;


Image::Image():
    m_width(0),
    m_height(0),
    m_components(0),
    m_exposure_time(0.0)
{
}

void Image::to_string(string &s) const {
    stringstream ss;
    ss << "Image(";
    ss << m_width << 'x' << m_height << 'x' << m_components << ")";
    s = ss.str();
}

bool Image::read_from_file(const char *file_path) {
    int x = 0;
    int y = 0;
    int components = 0;

    unsigned char *data = stbi_load(file_path, &x, &y, &components, 0);

    if (data) {
        m_width = x;
        m_height = y;
        m_components = components;
        size_t size = (size_t)x * (size_t)y * (size_t)components;
        m_image_data.assign(data, data + size);

        // try to read EXIF data
        {
            ExifTags tags;
            ifstream stream;
            stream.open(file_path);

            if (stream.is_open() and ExifTags::read_from_jpeg(stream, tags)) {
                if (tags.get_exposure_time(m_exposure_time)) {
                    cout << "Exposure time: " << m_exposure_time << " sec" << endl;
                }
            }
        }

        return true;
    } else {
        cerr << "Can not read \"" << file_path << "\" as an image." << endl;
        return false;
    }
    
    stbi_image_free(data);
}

struct my_error_mgr {
  struct jpeg_error_mgr pub;	/* "public" fields */

  jmp_buf setjmp_buffer;	/* for return to caller */
};

typedef struct my_error_mgr * my_error_ptr;

/*
 * Here's the routine that will replace the standard error_exit method:
 */

METHODDEF(void)
my_error_exit (j_common_ptr cinfo)
{
  /* cinfo->err really points to a my_error_mgr struct, so coerce pointer */
  my_error_ptr myerr = (my_error_ptr) cinfo->err;

  /* Always display the message. */
  /* We could postpone this until after returning, if we chose. */
  (*cinfo->err->output_message) (cinfo);

  /* Return control to the setjmp point */
  longjmp(myerr->setjmp_buffer, 1);
}

bool Image::read_from_jpeg_file(const char *file_path) {
    /* This struct contains the JPEG decompression parameters and pointers to
    * working space (which is allocated as needed by the JPEG library).
    */
    struct jpeg_decompress_struct cinfo;

    /* We use our private extension JPEG error handler.
    * Note that this struct must live as long as the main JPEG parameter
    * struct, to avoid dangling-pointer problems.
    */
    struct my_error_mgr jerr;

    /* More stuff */
    FILE * infile;		/* source file */
    JSAMPARRAY buffer;		/* Output row buffer */
    int row_stride;		/* physical row width in output buffer */

    /* In this example we want to open the input file before doing anything else,
    * so that the setjmp() error recovery below can assume the file is open.
    * VERY IMPORTANT: use "b" option to fopen() if you are on a machine that
    * requires it in order to read binary files.
    */

    if ((infile = fopen(file_path, "rb")) == NULL) {
        cerr << "Can not open \"" << file_path << "\"" << endl;
        return false;
    }

    /* Step 1: allocate and initialize JPEG decompression object */

    /* We set up the normal JPEG error routines, then override error_exit. */
    cinfo.err = jpeg_std_error(&jerr.pub);
    jerr.pub.error_exit = my_error_exit;
    /* Establish the setjmp return context for my_error_exit to use. */
    if (setjmp(jerr.setjmp_buffer)) {
        /* If we get here, the JPEG code has signaled an error.
         * We need to clean up the JPEG object, close the input file, and return.
         */
        jpeg_destroy_decompress(&cinfo);
        fclose(infile);
        return false;
    }
    /* Now we can initialize the JPEG decompression object. */
    jpeg_create_decompress(&cinfo);

    /* Step 2: specify data source (eg, a file) */

    jpeg_stdio_src(&cinfo, infile);

    /* Step 3: read file parameters with jpeg_read_header() */

    (void) jpeg_read_header(&cinfo, TRUE);
    /* We can ignore the return value from jpeg_read_header since
    *   (a) suspension is not possible with the stdio data source, and
    *   (b) we passed TRUE to reject a tables-only JPEG file as an error.
    * See libjpeg.txt for more info.
    */

    /* Step 4: set parameters for decompression */

    /* In this example, we don't need to change any of the defaults set by
    * jpeg_read_header(), so we do nothing here.
    */

    /* Step 5: Start decompressor */

    (void) jpeg_start_decompress(&cinfo);
    /* We can ignore the return value since suspension is not possible
    * with the stdio data source.
    */

    m_image_data.reserve(cinfo.output_width * cinfo.output_height * cinfo.output_components);

    m_width = cinfo.output_width;
    m_height = cinfo.output_height;
    m_components = cinfo.output_components;

    /* We may need to do some setup of our own at this point before reading
    * the data.  After jpeg_start_decompress() we have the correct scaled
    * output image dimensions available, as well as the output colormap
    * if we asked for color quantization.
    * In this example, we need to make an output work buffer of the right size.
    */ 
    /* JSAMPLEs per row in output buffer */
    row_stride = cinfo.output_width * cinfo.output_components;
    /* Make a one-row-high sample array that will go away when done with image */
    buffer = (*cinfo.mem->alloc_sarray)
        ((j_common_ptr) &cinfo, JPOOL_IMAGE, row_stride, 1);

    /* Step 6: while (scan lines remain to be read) */
    /*           jpeg_read_scanlines(...); */

    /* Here we use the library's state variable cinfo.output_scanline as the
    * loop counter, so that we don't have to keep track ourselves.
    */
    while (cinfo.output_scanline < cinfo.output_height) {
        /* jpeg_read_scanlines expects an array of pointers to scanlines.
         * Here the array is only one element long, but you could ask for
         * more than one scanline at a time if that's more convenient.
         */
        (void) jpeg_read_scanlines(&cinfo, buffer, 1);
        /* Assume put_scanline_someplace wants a pointer and sample count. */
        //put_scanline_someplace(buffer[0], row_stride);

        for (int i = 0; i < row_stride; i++) {
            m_image_data.push_back(buffer[0][i]);
        }
    }

    /* Step 7: Finish decompression */

    (void) jpeg_finish_decompress(&cinfo);
    /* We can ignore the return value since suspension is not possible
    * with the stdio data source.
    */

    /* Step 8: Release JPEG decompression object */

    /* This is an important step since it will release a good deal of memory. */
    jpeg_destroy_decompress(&cinfo);

    /* After finish_decompress, we can close the input file.
    * Here we postpone it until after no more JPEG errors are possible,
    * so as to simplify the setjmp error logic above.  (Actually, I don't
    * think that jpeg_destroy can do an error exit, but why assume anything...)
    */
    fclose(infile);

    if (jerr.pub.num_warnings != 0) {
        cout << jerr.pub.num_warnings << " errors occured." << endl;
        return false;
    } else {
        return true;
    }
}

bool Image::save_as_ppm_file(const char *file_path, bool ascii) const {
    if (ascii) {
        // ASCII

        ofstream out(file_path);
        out << "P3\n";
        out << m_width << ' ' << m_height << '\n';
        out << "255\n";

        unsigned int pixels = m_width * m_height;
        unsigned int index = 0;
        for (unsigned int i = 0; i < pixels; i++) {
            out << (int)m_image_data[index] << ' ' << (int)m_image_data[index + 1] << ' ' << (int)m_image_data[index + 2] << '\n';
            index += 3;
        }

        out.close();
        return true;
    } else {
        // binary

        ofstream out(file_path);
        out << "P6\n";
        out << m_width << ' ' << m_height << '\n';
        out << "255\n";

        unsigned int pixels = m_width * m_height;
        unsigned int index = 0;
        for (unsigned int i = 0; i < pixels; i++) {
            out << m_image_data[index] << m_image_data[index + 1] << m_image_data[index + 2];
            index += 3;
        }

        out.close();
        return true;
    }
}

unsigned int Image::get_width() const {
    return m_width;
}

unsigned int Image::get_height() const {
    return m_height;
}

unsigned int Image::get_components() const {
    return m_components;
}

bool Image::has_equal_dimensions(const Image &other) const {
    return (m_width == other.m_width &&
            m_height == other.m_height &&
            m_components == other.m_components);
}

double Image::get_exposure_time() const {
    return m_exposure_time;
}