/*
 * OMAP3 ISP snapshot test application
 *
 * Copyright (C) 2012 Ideas on board SPRL
 *
 * Contact: Laurent Pinchart <laurent.pinchart@ideasonboard.com>
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the Free
 * Software Foundation; either version 2 of the License, or (at your option)
 * any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program; if not, write to the Free Software Foundation, Inc.,
 * 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 */

#define _BSD_SOURCE
#include <errno.h>
#include <setjmp.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <jpeglib.h>

#include "isp/tools.h"
#include "isp/v4l2-pool.h"

#include "jpeg.h"

struct jpeg {
	struct jpeg_compress_struct cinfo;
	struct jpeg_error_mgr error;
	unsigned int width;
	unsigned int height;
	uint8_t *buffer;
	jmp_buf jump;
};

static void jpeg_error_exit(j_common_ptr cinfo)
{
	struct jpeg *jpeg = container_of(cinfo, struct jpeg, cinfo);

	longjmp(jpeg->jump, 1);
}

int jpeg_save(struct jpeg *jpeg, FILE *file, struct v4l2_video_buffer *buffer)
{
	JSAMPROW rows_y[DCTSIZE];
	JSAMPROW rows_u[DCTSIZE];
	JSAMPROW rows_v[DCTSIZE];
	JSAMPARRAY image[3];
	unsigned int width;
	unsigned int height;
	uint8_t *planes[3];
	unsigned int x;
	unsigned int y;
	uint8_t *mem;
	unsigned int i;

	/* Convert the YUYV input buffer to a planar format for JPEG
	 * compression.
	 */
	width = (jpeg->width + 2 * DCTSIZE - 1) / (2 * DCTSIZE) * (2 * DCTSIZE);
	height = (jpeg->height + DCTSIZE - 1) / DCTSIZE * DCTSIZE;

	planes[0] = jpeg->buffer;
	planes[1] = jpeg->buffer + width * height;
	planes[2] = jpeg->buffer + 3 * width / 2 * height;

	mem = buffer->mem;

	for (y = 0; y < jpeg->height; ++y) {
		for (x = 0; x < jpeg->width; x += 2) {
			planes[0][x] = *mem++;
			planes[1][x/2] = *mem++;
			planes[0][x+1] = *mem++;
			planes[2][x/2] = *mem++;
		}
		planes[0] += width;
		planes[1] += width / 2;
		planes[2] += width / 2;
	}

	/* Now perform JPEG compression. */
	if (setjmp(jpeg->jump))
		return -EIO;

	jpeg_stdio_dest(&jpeg->cinfo, file);
	jpeg_start_compress(&jpeg->cinfo, TRUE);

	rows_y[0] = jpeg->buffer;
	rows_u[0] = jpeg->buffer + width * height;
	rows_v[0] = jpeg->buffer + 3 * width / 2 * height;

	for (i = 1; i < DCTSIZE; ++i) {
		rows_y[i] = rows_y[i-1] + width;
		rows_u[i] = rows_u[i-1] + width / 2;
		rows_v[i] = rows_v[i-1] + width / 2;
	}

	image[0] = rows_y;
	image[1] = rows_u;
	image[2] = rows_v;

	while (jpeg->cinfo.next_scanline < jpeg->height) {
		unsigned int written;

		written = jpeg_write_raw_data(&jpeg->cinfo, image, DCTSIZE);

		for (i = 0; i < DCTSIZE; ++i) {
			rows_y[i] += width * written;
			rows_u[i] += width / 2 * written;
			rows_v[i] += width / 2 * written;
		}
	}

	jpeg_finish_compress(&jpeg->cinfo);

	return 0;
}

static int jpeg_init_compressor(struct jpeg *jpeg)
{
	jpeg->cinfo.err = jpeg_std_error(&jpeg->error);
	jpeg->error.error_exit = jpeg_error_exit;
	if (setjmp(jpeg->jump)) {
		jpeg_destroy_compress(&jpeg->cinfo);
		jpeg->cinfo.err = NULL;
		printf("Unable to create JPEG compression handler.\n");
		return -EIO;
	}

	jpeg_create_compress(&jpeg->cinfo);

	jpeg->cinfo.image_width = jpeg->width;
	jpeg->cinfo.image_height = jpeg->height;
	jpeg->cinfo.input_components = 3;
	jpeg->cinfo.in_color_space = JCS_YCbCr;

	jpeg_set_defaults(&jpeg->cinfo);
	jpeg_set_quality(&jpeg->cinfo, 80, TRUE);

	jpeg->cinfo.raw_data_in = TRUE;
	jpeg_set_colorspace(&jpeg->cinfo, JCS_YCbCr);
	jpeg->cinfo.comp_info[0].h_samp_factor = 2;
	jpeg->cinfo.comp_info[0].v_samp_factor = 1;
	jpeg->cinfo.comp_info[1].h_samp_factor = 1;
	jpeg->cinfo.comp_info[1].v_samp_factor = 1;
	jpeg->cinfo.comp_info[2].h_samp_factor = 1;
	jpeg->cinfo.comp_info[2].v_samp_factor = 1;

	return 0;
}

struct jpeg *jpeg_init(const struct v4l2_mbus_framefmt *format)
{
	unsigned int width;
	unsigned int height;
	struct jpeg *jpeg;

	jpeg = malloc(sizeof *jpeg);
	if (jpeg == NULL)
		return NULL;

	memset(jpeg, 0, sizeof *jpeg);
	jpeg->width = format->width;
	jpeg->height = format->height;

	if (jpeg_init_compressor(jpeg) < 0) {
		free(jpeg);
		return NULL;
	}

	/* Allocate a temporary buffer to hold the planar YUV image.
	 *
	 * All planes must be padded to the DCT block size horizontally
	 * and vertically. The luma plane must thus be padded to twice
	 * the DCT block size horizontally as the chroma planes are
	 * subsampled by 2 in the horizontal direction. For simplicity,
	 * pad all planes to twice the DCT block size horizontally.
	 */
	width = (format->width + 2 * DCTSIZE - 1) / (2 * DCTSIZE) * (2 * DCTSIZE);
	height = (format->height + DCTSIZE - 1) / DCTSIZE * DCTSIZE;

	jpeg->buffer = malloc(width * height * 2);
	if (jpeg->buffer == NULL) {
		jpeg_destroy_compress(&jpeg->cinfo);
		jpeg->cinfo.err = NULL;
		free(jpeg);
		printf("Unable to allocate memory for JPEG compressor.\n");
		return NULL;
	}

	return jpeg;
}

void jpeg_cleanup(struct jpeg *jpeg)
{
	if (jpeg->cinfo.err != NULL) {
		jpeg_destroy_compress(&jpeg->cinfo);
		jpeg->cinfo.err = NULL;
	}
	free(jpeg->buffer);
	free(jpeg);
}