#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <jpeglib.h>
#include <omp.h>

static int blur_radius = 15;
static int blur_times = 5;

static void blur_h(
		unsigned char *dest, unsigned char *src,
		int stride, int width, int height, int radius) {
	double coeff = 1.0 / (radius * 2 + 1);

#pragma omp parallel for
	for (int i = 0; i < height; ++i) {
		int iwidth = i * stride * 3;
		double r_acc = 0.0;
		double g_acc = 0.0;
		double b_acc = 0.0;
		for (int j = -radius; j < width; ++j) {
			if (j - radius - 1 >= 0) {
				int index = iwidth + (j - radius - 1) * 3;
				r_acc -= coeff * src[index + 0];
				g_acc -= coeff * src[index + 1];
				b_acc -= coeff * src[index + 2];
			}
			if (j + radius < width) {
				int index = iwidth + (j + radius) * 3;
				r_acc += coeff * src[index + 0];
				g_acc += coeff * src[index + 1];
				b_acc += coeff * src[index + 2];
			}
			if (j < 0)
				continue;
			int index = iwidth + j * 3;
			dest[index + 0] = r_acc + 0.5;
			dest[index + 1] = g_acc + 0.5;
			dest[index + 2] = b_acc + 0.5;
		}
	}
}

static void blur_v(
		unsigned char *dest, unsigned char *src,
		int stride, int width, int height, int radius) {
	double coeff = 1.0 / (radius * 2 + 1);

#pragma omp parallel for
	for (int j = 0; j < width; ++j) {
		double r_acc = 0.0;
		double g_acc = 0.0;
		double b_acc = 0.0;
		for (int i = -radius; i < height; ++i) {
			if (i - radius - 1 >= 0) {
				int index = (i - radius - 1) * stride * 3 + j * 3;
				r_acc -= coeff * src[index + 0];
				g_acc -= coeff * src[index + 1];
				b_acc -= coeff * src[index + 2];
			}
			if (i + radius < height) {
				int index = (i + radius) * stride * 3 + j * 3;
				r_acc += coeff * src[index + 0];
				g_acc += coeff * src[index + 1];
				b_acc += coeff * src[index + 2];
			}
			if (i < 0)
				continue;
			int index = i * stride * 3 + j * 3;
			dest[index + 0] = r_acc + 0.5;
			dest[index + 1] = g_acc + 0.5;
			dest[index + 2] = b_acc + 0.5;
		}
	}
}

static void blur_once(
		unsigned char *dest, unsigned char *src, unsigned char *scratch,
		int stride, int width, int height, int radius) {
	blur_h(scratch, src, stride, width, height, radius);
	blur_v(dest, scratch, stride, width, height, radius);
}

void blur_rect(
		unsigned char *image, unsigned char *scratch, unsigned char *scratch2,
		int stride, int x, int y, int w, int h, int radius, int times) {
	if (w == 0 || h == 0)
		return;

	unsigned char *orig = image;
	int idx = y * stride * 3 + x * 3;
	for (int i = 0; i < times; ++i) {
		blur_once(scratch + idx, image + idx, scratch2, stride, w, h, radius);
		if (i != times - 1) {
			unsigned char *tmp = image;
			image = scratch;
			scratch = tmp;
		}
	}

	if (image != orig)
		memcpy(scratch + idx, image + idx, stride * h * 3);
}

unsigned char *process_image(
		unsigned char *input_image, int input_width, int input_height,
		int desired_width, int desired_height) {
	unsigned char *output_image = calloc(3, desired_width * desired_height);

	double scale_w = (double)desired_width / (double)input_width;
	double scale_h = (double)desired_height / (double)input_height;
	double scale;

	int rect_x, rect_y, rect_w, rect_h;
	if (scale_w < scale_h)
		scale = scale_w;
	else
		scale = scale_h;

	rect_w = input_width * scale;
	rect_h = input_height * scale;
	rect_x = (desired_width - rect_w) / 2.0;
	rect_y = (desired_height - rect_h) / 2.0;


	// Write main image
	for (int y = 0; y < rect_h; ++y) {
		for (int x = 0; x < rect_w; ++x) {
			int input_y = y / scale;
			int input_x = x / scale;
			int input_idx = input_y * input_width * 3 + input_x * 3;
			int output_idx = (y + rect_y) * desired_width * 3 + (x + rect_x) * 3;
			output_image[output_idx + 0] = input_image[input_idx + 0];
			output_image[output_idx + 1] = input_image[input_idx + 1];
			output_image[output_idx + 2] = input_image[input_idx + 2];
		}
	}

	// Top/bottom
	for (int y = 0; y < rect_y; ++y) {
		for (int x = 0; x < rect_w; ++x) {
			int input_idx = (y + rect_y) * desired_width * 3 + (x + rect_x) * 3;
			int output_idx = y * desired_width * 3 + (x + rect_x) * 3;
			output_image[output_idx + 0] = output_image[input_idx + 0];
			output_image[output_idx + 1] = output_image[input_idx + 1];
			output_image[output_idx + 2] = output_image[input_idx + 2];

			input_idx = (y + rect_h) * desired_width * 3 + (x + rect_x) * 3;
			output_idx = (y + rect_h + rect_y) * desired_width * 3 + (x + rect_x) * 3;
			output_image[output_idx + 0] = output_image[input_idx + 0];
			output_image[output_idx + 1] = output_image[input_idx + 1];
			output_image[output_idx + 2] = output_image[input_idx + 2];
		}
	}

	// Left/right
	for (int x = 0; x < rect_x; ++x) {
		for (int y = 0; y < rect_h; ++y) {
			int input_idx_l = (y + rect_y) * desired_width * 3 + (x + rect_x) * 3;
			int output_idx_l = y * desired_width * 3 + (x) * 3;
			int input_idx_r = (y + rect_y) * desired_width * 3 + (x + rect_w) * 3;
			int output_idx_r = y * desired_width * 3 + (x + rect_w + rect_x) * 3;

			output_image[output_idx_l + 0] = output_image[input_idx_l + 0];
			output_image[output_idx_l + 1] = output_image[input_idx_l + 1];
			output_image[output_idx_l + 2] = output_image[input_idx_l + 2];

			output_image[output_idx_r + 0] = output_image[input_idx_r + 0];
			output_image[output_idx_r + 1] = output_image[input_idx_r + 1];
			output_image[output_idx_r + 2] = output_image[input_idx_r + 2];
		}
	}

	// Blur
	unsigned char *scratch = calloc(3, desired_width * desired_height);
	unsigned char *scratch2 = calloc(3, desired_width * desired_height);
	blur_rect(
			output_image, scratch, scratch2, desired_width,
			0, 0, desired_width, rect_y, blur_radius, blur_times);
	blur_rect(
			output_image, scratch, scratch2, desired_width,
			0, rect_y + rect_h, desired_width, rect_y, blur_radius, blur_times);
	blur_rect(
			output_image, scratch, scratch2, desired_width,
			0, 0, rect_x, desired_height, blur_radius, blur_times);
	blur_rect(
			output_image, scratch, scratch2, desired_width,
			rect_x + rect_w, 0, rect_x, desired_height, blur_radius, blur_times);
	free(scratch);
	free(scratch2);

	return output_image;
}

unsigned char *read_jpg(FILE *in, int *width, int *height) {
	struct jpeg_decompress_struct cinfo = { 0 };
	struct jpeg_error_mgr jerr = { 0 };
	cinfo.err = jpeg_std_error(&jerr);

	jpeg_create_decompress(&cinfo);
	jpeg_stdio_src(&cinfo, in);
	jpeg_read_header(&cinfo, TRUE);
	jpeg_start_decompress(&cinfo);

	*width = cinfo.output_width;
	*height = cinfo.output_height;
	int bytes_per_line = cinfo.output_width * 3;
	unsigned char *image = malloc(bytes_per_line * cinfo.output_height);

	/* Make a one-row-high sample array that will go away when done with image */
	JSAMPARRAY buffer = (*cinfo.mem->alloc_sarray)
		((j_common_ptr)&cinfo, JPOOL_IMAGE, bytes_per_line, 1);

	while (cinfo.output_scanline < cinfo.output_height) {
		jpeg_read_scanlines(&cinfo, buffer, 1);
		/* Assume put_scanline_someplace wants a pointer and sample count. */
		memcpy(&image[bytes_per_line * (cinfo.output_scanline - 1)], buffer[0], bytes_per_line);
	}

	jpeg_finish_decompress(&cinfo);
	jpeg_destroy_decompress(&cinfo);

	return image;
}

void write_jpg(FILE *out, unsigned char *image, int width, int height) {
	struct jpeg_compress_struct cinfo = { 0 };
	struct jpeg_error_mgr jerr = { 0 };
	cinfo.err = jpeg_std_error(&jerr);

	jpeg_create_compress(&cinfo);
	jpeg_stdio_dest(&cinfo, out);

	cinfo.image_width = width - blur_radius * 2;
	cinfo.image_height = height - blur_radius * 2;
	cinfo.input_components = 3;
	cinfo.in_color_space = JCS_RGB;

	jpeg_set_defaults(&cinfo);
	jpeg_set_quality(&cinfo, 90, FALSE);
	jpeg_start_compress(&cinfo, TRUE);

	int bytes_per_line = width * 3;
	JSAMPROW row_pointer[1];

	while (cinfo.next_scanline < cinfo.image_height) {
		row_pointer[0] = &image[(cinfo.next_scanline + blur_radius) * bytes_per_line + blur_radius * 3];
		jpeg_write_scanlines(&cinfo, row_pointer, 1);
	}

	jpeg_finish_compress(&cinfo);
	jpeg_destroy_compress(&cinfo);
}

int main(int argc, char **argv) {
	if (argc != 4) {
		printf("Usage: %s <resolution> <input> <output>\n", argv[0]);
		return EXIT_FAILURE;
	}

	int desired_width, desired_height;
	if (sscanf(argv[1], "%ix%i", &desired_width, &desired_height) != 2) {
		fprintf(stderr, "Error parsing resolution '%s'\n", argv[1]);
		return EXIT_FAILURE;
	}
	desired_width += blur_radius * 2;
	desired_height += blur_radius * 2;

	FILE *in;
	if (strcmp(argv[2], "-") == 0) {
		in = stdin;
	} else {
		in = fopen(argv[2], "r");
		if (in == NULL) {
			perror(argv[2]);
			return EXIT_FAILURE;
		}
	}

	int input_width, input_height;
	unsigned char *input_image = read_jpg(in, &input_width, &input_height);
	fclose(in);

	unsigned char *output_image = process_image(
			input_image, input_width, input_height, desired_width, desired_height);
	free(input_image);

	FILE *out;
	if (strcmp(argv[3], "-") == 0) {
		out = stdout;
	} else {
		out = fopen(argv[3], "w");
		if (out == NULL) {
			perror(argv[3]);
			free(output_image);
			return EXIT_FAILURE;
		}
	}

	write_jpg(out, output_image, desired_width, desired_height);
	free(output_image);

	return EXIT_SUCCESS;
}