/*
    Scan Tailor - Interactive post-processing tool for scanned pages.
	Copyright (C) 2007-2009  Joseph Artsimovich <[email protected]>

    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 3 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, see <http://www.gnu.org/licenses/>.
*/

#include "JpegMetadataLoader.h"
#include "ImageMetadata.h"
#include "NonCopyable.h"
#include "Dpi.h"
#include "Dpm.h"
#include <QIODevice>
#include <QSize>
#include <QDebug>
#include <setjmp.h>
#include <string.h>
#include <assert.h>

extern "C" {
#include <jpeglib.h>
}

namespace
{

/*======================== JpegDecompressionHandle =======================*/

class JpegDecompressHandle
{
	DECLARE_NON_COPYABLE(JpegDecompressHandle)
public:
	JpegDecompressHandle(jpeg_error_mgr* err_mgr, jpeg_source_mgr* src_mgr);
	
	~JpegDecompressHandle();
	
	jpeg_decompress_struct* ptr() { return &m_info; }
	
	jpeg_decompress_struct* operator->() { return &m_info; }
private:
	jpeg_decompress_struct m_info;
};

JpegDecompressHandle::JpegDecompressHandle(
	jpeg_error_mgr* err_mgr, jpeg_source_mgr* src_mgr)
{
	m_info.err = err_mgr;
	jpeg_create_decompress(&m_info);
	m_info.src = src_mgr;
}

JpegDecompressHandle::~JpegDecompressHandle()
{
	jpeg_destroy_decompress(&m_info);
}


/*============================ JpegSourceManager =========================*/

class JpegSourceManager : public jpeg_source_mgr
{
	DECLARE_NON_COPYABLE(JpegSourceManager)
public:
	JpegSourceManager(QIODevice& io_device);
private:
	static void initSource(j_decompress_ptr cinfo);
	
	static boolean fillInputBuffer(j_decompress_ptr cinfo);
	
	boolean fillInputBufferImpl();
	
	static void skipInputData(j_decompress_ptr cinfo, long num_bytes);
	
	void skipInputDataImpl(long num_bytes);
	
	static void termSource(j_decompress_ptr cinfo);
	
	static JpegSourceManager* object(j_decompress_ptr cinfo);
	
	QIODevice& m_rDevice;
	JOCTET m_buf[4096];
};

JpegSourceManager::JpegSourceManager(QIODevice& io_device)
:	m_rDevice(io_device)
{
	init_source = &JpegSourceManager::initSource;
	fill_input_buffer = &JpegSourceManager::fillInputBuffer;
	skip_input_data = &JpegSourceManager::skipInputData;
	resync_to_restart = &jpeg_resync_to_restart;
	term_source = &JpegSourceManager::termSource;
	bytes_in_buffer = 0;
	next_input_byte = m_buf;
}

void
JpegSourceManager::initSource(j_decompress_ptr cinfo)
{
	// No-op.
}

boolean
JpegSourceManager::fillInputBuffer(j_decompress_ptr cinfo)
{
	return object(cinfo)->fillInputBufferImpl();
}

boolean
JpegSourceManager::fillInputBufferImpl()
{
	qint64 const bytes_read = m_rDevice.read((char*)m_buf, sizeof(m_buf));
	if (bytes_read > 0) {
		bytes_in_buffer = bytes_read;
	} else {
		// Insert a fake EOI marker.
		m_buf[0] = 0xFF;
		m_buf[1] = JPEG_EOI;
		bytes_in_buffer = 2;
	}
	next_input_byte = m_buf;
	return 1;
}

void
JpegSourceManager::skipInputData(j_decompress_ptr cinfo, long num_bytes)
{
	object(cinfo)->skipInputDataImpl(num_bytes);
}

void
JpegSourceManager::skipInputDataImpl(long num_bytes)
{
	if (num_bytes <= 0) {
		return;
	}
	
	while (num_bytes > (long)bytes_in_buffer) {
		num_bytes -= (long)bytes_in_buffer;
		fillInputBufferImpl();
	}
	next_input_byte += num_bytes;
	bytes_in_buffer -= num_bytes;
}

void
JpegSourceManager::termSource(j_decompress_ptr cinfo)
{
	// No-op.
}

JpegSourceManager*
JpegSourceManager::object(j_decompress_ptr cinfo)
{
	return static_cast<JpegSourceManager*>(cinfo->src);
}


/*============================= JpegErrorManager ===========================*/

class JpegErrorManager : public jpeg_error_mgr
{
	DECLARE_NON_COPYABLE(JpegErrorManager)
public:
	JpegErrorManager();
	
	jmp_buf& jmpBuf() { return m_jmpBuf; }
private:
	static void errorExit(j_common_ptr cinfo);
	
	static JpegErrorManager* object(j_common_ptr cinfo);
	
	jmp_buf m_jmpBuf;
};

JpegErrorManager::JpegErrorManager()
{
	jpeg_std_error(this);
	error_exit = &JpegErrorManager::errorExit;
}

void
JpegErrorManager::errorExit(j_common_ptr cinfo)
{
	longjmp(object(cinfo)->jmpBuf(), 1);
}

JpegErrorManager*
JpegErrorManager::object(j_common_ptr cinfo)
{
	return static_cast<JpegErrorManager*>(cinfo->err);
}

} // anonymous namespace


/*============================= JpegMetadataLoader ==========================*/

void
JpegMetadataLoader::registerMyself()
{
	static bool registered = false;
	if (!registered) {
		ImageMetadataLoader::registerLoader(
			IntrusivePtr<ImageMetadataLoader>(new JpegMetadataLoader)
		);
		registered = true;
	}
}

ImageMetadataLoader::Status
JpegMetadataLoader::loadMetadata(
	QIODevice& io_device,
	VirtualFunction1<void, ImageMetadata const&>& out)
{
	if (!io_device.isReadable()) {
		return GENERIC_ERROR;
	}

	static unsigned char const jpeg_signature[] = { 0xff, 0xd8, 0xff };
	static int const sig_size = sizeof(jpeg_signature);

	unsigned char signature[sig_size];
	if (io_device.peek((char*)signature, sig_size) != sig_size) {
		return FORMAT_NOT_RECOGNIZED;
	}
	if (memcmp(jpeg_signature, signature, sig_size) != 0) {
		return FORMAT_NOT_RECOGNIZED;
	}
	
	JpegErrorManager err_mgr;
	if (setjmp(err_mgr.jmpBuf())) {
		// Returning from longjmp().
		return GENERIC_ERROR;
	}
	
	JpegSourceManager src_mgr(io_device);
	JpegDecompressHandle cinfo(&err_mgr, &src_mgr);
	
	int const header_status = jpeg_read_header(cinfo.ptr(), 0);
	if (header_status == JPEG_HEADER_TABLES_ONLY) {
		return NO_IMAGES;
	}
	
	// The other possible value is JPEG_SUSPENDED, but we never suspend it.
	assert(header_status == JPEG_HEADER_OK);
	
	if (!jpeg_start_decompress(cinfo.ptr())) {
		// libjpeg doesn't support all compression types.
		return GENERIC_ERROR;
	}
	
	QSize const size(cinfo->image_width, cinfo->image_height);
	Dpi dpi;
	if (cinfo->density_unit == 1) {
		// Dots per inch.
		dpi = Dpi(cinfo->X_density, cinfo->Y_density);
	} else if (cinfo->density_unit == 2) {
		// Dots per centimeter.
		dpi = Dpm(cinfo->X_density * 100, cinfo->Y_density * 100);
	}
	
	out(ImageMetadata(size, dpi));
	return LOADED;
}