#include <stdio.h>
#include <stdlib.h>
#include <sys/types.h>
#include <jpeglib.h>
#include <string.h>
#include <setjmp.h>
#include <sys/stat.h>
#include <math.h>
#include "pqstego.h"


static long jround_up(long a, long b) {
    a += b - 1L;
    return a - (a % b);
}

struct pqstego_jpeg_error_mgr {
    struct jpeg_error_mgr pub;
    jmp_buf setjmp_buffer;
};

typedef struct pqstego_jpeg_error_mgr * pqstego_jpeg_error_ptr;

void pqstego_jpeg_error_exit(j_common_ptr cinfo) {
    pqstego_jpeg_error_ptr err = (pqstego_jpeg_error_ptr) cinfo->err;
    (*cinfo->err->output_message) (cinfo);
    longjmp(err->setjmp_buffer, 1);
}


uint8_t main(int argc, char **argv) {
    
    uint8_t *coverfilename = "cover.jpg";
    uint8_t *stegofilename = "stego.jpg";
    pq_parameter_t *param = NULL;
    struct pqstego_jpeg_error_mgr jerr;
    jvirt_barray_ptr *coef_array;
    JBLOCKARRAY coef_buffer;
    JDIMENSION i, ii, iii;
    struct jpeg_decompress_struct *dinfo;
    pq_data_t *pq_src_data = NULL;
    pq_data_t *pq_dst_data = NULL;
    uint8_t *message = "HelloWorld";
    uint32_t msglen = 11;
    uint8_t exit = EXIT_SUCCESS;

    //Copy JPEG quantization table and coefficients to pq_data_t struct
    pq_src_data = (pq_data_t*) malloc(sizeof (pq_data_t));
    pq_dst_data = (pq_data_t*) malloc(sizeof (pq_data_t));
    dinfo = malloc(sizeof (struct jpeg_decompress_struct));
    dinfo->err = jpeg_std_error(&jerr.pub);

    if ((srcfile = fopen(coverfilename, "rb")) == NULL) {
        return EXIT_FAILURE;
    }
    jpeg_create_decompress(dinfo);
    jpeg_stdio_src(dinfo, srcfile);
    jpeg_read_header(dinfo, TRUE);

    coef_array = jpeg_read_coefficients(dinfo);
    pq_src_data->components = dinfo->num_components;
    pq_src_data->component = malloc(sizeof (pq_component_t) * dinfo->num_components);
    if (!pq_src_data->component) {
        return EXIT_FAILURE;
    }
    pq_src_data->size_x = dinfo->image_width;
    pq_src_data->size_y = dinfo->image_height;


    for (i = 0; i < dinfo->num_components; i++) {
        jpeg_component_info *compinfo = (dinfo->comp_info + i);
        JDIMENSION blocks = compinfo->width_in_blocks * compinfo->height_in_blocks;
        (pq_src_data->component + i)->quant = (pq_quant_t*) malloc(sizeof (pq_quant_t));
        if (!(pq_src_data->component + i)->quant) {
            return EXIT_FAILURE;
        }
        (pq_src_data->component + i)->block = (pq_block_t*) malloc(sizeof (pq_block_t) * blocks);
        if (!(pq_src_data->component + i)->block) {
            return EXIT_FAILURE;
        }
        (pq_src_data->component + i)->blocks = blocks;
        memcpy((pq_src_data->component + i)->quant->values, compinfo->quant_table->quantval, sizeof (pq_quant_t));
        for (ii = 0; ii < compinfo->height_in_blocks; ii++) {
            coef_buffer = (dinfo->mem->access_virt_barray)
                    ((j_common_ptr) dinfo, *(coef_array + i), ii, (JDIMENSION) 1, FALSE);
            JBLOCKROW block_row = *(coef_buffer);
            for (iii = 0; iii < compinfo->width_in_blocks; iii++) {
                memcpy(((pq_src_data->component + i)->block + (ii * compinfo->width_in_blocks) + iii)->values,
                        (block_row + iii), sizeof (pq_block_t));
            }

        }

    }
    fclose(srcfile);

    //Allocate pq_parameter_t struct
    param = (pq_parameter_t*) malloc(sizeof (pq_parameter_t));
    
    //Password
    param->password = "password";
    //Length of the password
    param->pwlen = 9;
    //Set this to more than 0 if you wanna force quality of 
    //the stego image, else the algorithm will take the quality 
    //with the highest capacity 
    param->quality = 0;
    //Start embed, the resulting image will be in pq_dst_data
    pq_embed(pq_src_data, pq_dst_data, message, msglen, param);
    //Extract the message again from pq_dst_data
    pq_extract(pq_dst_data, &message, &msglen, param);

    return EXIT_SUCCESS;
}


