#include <QtCore/QtEndian>
#include <QtCore/QBuffer>
#include <QtCore/QVariant>
#include <QtCore/QDebug>
#include <QtGui/QImage>

#include "msolimageiohandler.h"

static const int header_size = 16;

namespace
{

static inline uchar encode_color(QRgb c)
{
	return (qRed(c) & 0xC0) | ((qGreen(c) & 0xC0) >> 2) | ((qBlue(c) & 0xC0) >> 4) | 3;
}

static inline QRgb decode_color(uchar p, uchar a = 255)
{
	return qRgba(p & 0xC0, (p & 0x30) << 2, (p & 0x0C) << 4, a);
}

static inline uchar add_saturate(uchar val, int addend)
{
	int sum = val + addend;
	return qMax(0, qMin(sum, 255));
}

static inline void add_saturate(uchar* val, int addend)
{
	*val = add_saturate(*val, addend);
}

}

MSOLImageIOHandler::MSOLImageIOHandler() : QImageIOHandler()
{
}

bool MSOLImageIOHandler::canRead() const
{
	char header[header_size];
	if (device()->peek(header, header_size) != header_size) {
		return false;
	}

	return strncmp(&header[0], "MSOL  ", 6) == 0;
}

bool MSOLImageIOHandler::read(QImage *image)
{
	QIODevice *dev = device();

	uchar header[header_size];
	if (dev->read(reinterpret_cast<char*>(header), header_size) != header_size) {
		return false;
	}

	if (strncmp(reinterpret_cast<char*>(&header[0]), "MSOL  ", 6) != 0) {
		return false;
	}

	const bool alpha = header[6];
	const int width = qFromLittleEndian<quint16>(&header[8]);
	const int height = qFromLittleEndian<quint16>(&header[10]);
	const int pixel_size = alpha ? 2 : 1;
	const int line_size = width * pixel_size;

	if (alpha) {
		*image = QImage(width, height, QImage::Format_ARGB32);
	} else {
		*image = QImage(width, height, QImage::Format_RGB32);
	}

	Q_ASSERT(image->depth() == sizeof(QRgb) * 8);

	for (int y = 0; y < height; y++) {
		QRgb *dst = reinterpret_cast<QRgb*>(image->scanLine(y));
		uchar src[line_size];
		if (dev->read(reinterpret_cast<char*>(src), line_size) != line_size) {
			qDebug() << "Could not read image line" << y;
			return false;
		}
		for (int x = 0; x < width; x++) {
			if (alpha) {
				dst[x] = decode_color(src[x * pixel_size]);
			} else {
				dst[x] = decode_color(src[x * pixel_size], src[(x * pixel_size) + 1]);
			}
		}
	}

	return true;
}

bool MSOLImageIOHandler::write(const QImage &image)
{
	QIODevice *dev = device();
	const int height = image.height();
	const int width = image.width();
	const bool alpha = image.hasAlphaChannel();
	const int pixel_size = alpha ? 2 : 1;
	const int line_size = width * pixel_size;

	QImage img = image.convertToFormat(alpha ? QImage::Format_ARGB32 : QImage::Format_RGB32);
	img = ditherForMsol(img);

	uchar header[header_size] = { 0 };
	strcpy(reinterpret_cast<char*>(&header[0]), "MSOL  ");
	if (alpha) header[6] = 0x80;
	header[7] = 8;

	qToLittleEndian<quint16>(width, &header[8]);
	qToLittleEndian<quint16>(height, &header[10]);

	if (dev->write(reinterpret_cast<char*>(header), header_size) != header_size) {
		qDebug() << "Could not write header";
		return false;
	}

	Q_ASSERT(img.depth() == sizeof(QRgb) * 8);

	for (int y = 0; y < height; y++) {
		uchar dst[line_size];
		const QRgb *src = reinterpret_cast<const QRgb*>(img.constScanLine(y));
		for (int x = 0; x < width; x++) {
			dst[x * pixel_size] = encode_color(src[x]);
			if (alpha) {
				dst[(x * pixel_size) + 1] = qAlpha(src[x]);
			}
		}
		if (dev->write(reinterpret_cast<char*>(dst), line_size) != line_size) {
			qDebug() << "Could not write image line " << y;
			return false;
		}
	}

	return true;
}

bool MSOLImageIOHandler::supportsOption(ImageOption option) const
{
	switch (option) {
	case Size:
		return true;
	default:
		return false;
	}
}

QVariant MSOLImageIOHandler::option(ImageOption option) const
{
	char header[header_size];

	switch (option) {
	case Size:
		if (device()->peek(header, header_size) == header_size) {
			const QSize size(qFromLittleEndian<quint16>(reinterpret_cast<uchar*>(&header[8])),
							 qFromLittleEndian<quint16>(reinterpret_cast<uchar*>(&header[10])));
			return QVariant::fromValue(size);
		} else {
			qDebug() << "Cannot read from device: " << device()->errorString();
		}
	default:
		break;
	}

	return QVariant();
}

void MSOLImageIOHandler::setOption(ImageOption option, const QVariant &value)
{
	Q_UNUSED(option);
	Q_UNUSED(value);
}

QByteArray convertImageToMsol(const QImage &img)
{
	MSOLImageIOHandler msolio;
	QBuffer msolImg;
	msolImg.open(QBuffer::WriteOnly);
	msolio.setDevice(&msolImg);
	if (!msolio.write(img)) {
		return QByteArray();
	}
	return msolImg.data();
}

QImage convertMsolToImage(const QByteArray &msol)
{
	MSOLImageIOHandler msolio;
	QBuffer msolImg;
	QImage img;
	msolImg.setData(msol);
	msolImg.open(QBuffer::ReadOnly);
	if (!msolio.read(&img)) {
		return QImage();
	}
	return img;
}

QImage ditherForMsol(const QImage &src)
{
	Q_ASSERT(src.format() == QImage::Format_ARGB32 || src.format() == QImage::Format_RGB32);
	const int width = src.width(), height = src.height();
	const int bytesPerLine = src.bytesPerLine();

	QImage dst(src);

	// If we are on Big endian, then we must skip the first channel (=alpha).
	// In LE, use the first three channels and skip the final one (=alpha).
	const bool bigEndian = QSysInfo::ByteOrder == QSysInfo::BigEndian;
	uchar *line = dst.bits() + (bigEndian ? 1 : 0);
	for (int y = 0; y < height; y++) {
		for (int x = 0; x < width; x++) {
			for (int chan = 0; chan < 3; chan++) {
				const uchar val = line[x * 4 + chan];
				const uchar quant = add_saturate(val, 0x20) & 0xC0;
				const int error = val - quant;

				line[x * 4 + chan] = quant;

				if (x + 1 < width)
					add_saturate(&line[(x + 1) * 4 + chan], (error * 7) / 16);

				if (y + 1 < height) {
					uchar *nextLine = line + bytesPerLine;

					if (x > 0)
						add_saturate(&nextLine[(x - 1) * 4 + chan], (error * 3) / 16);

					add_saturate(&nextLine[x * 4 + chan], (error * 5) / 16);

					if (x + 1 < width)
						add_saturate(&nextLine[(x + 1) * 4 + chan], (error * 1) / 16);
				}
			}
		}

		line += bytesPerLine;
	}

	return dst;
}