#include "HGPnm.h"
#include "../base/HGInc.h"
#include "../base/HGInfo.h"
#include <vector>
#include <string>

static inline HGByte GetBit(const HGByte* data, HGUInt index)
{
	HGUInt byteIndex = index / 8;
	HGUInt bitIndex = index % 8;
	return (data[byteIndex] >> (7 - bitIndex)) & 0x01;
}

static inline void SetBit(HGByte* data, HGUInt index, HGByte value)
{
	assert(0 == value || 1 == value);
	HGUInt byteIndex = index / 8;
	HGUInt bitIndex = index % 8;
	if (1 == value)
		data[byteIndex] |= (1 << (7 - bitIndex));
	else
		data[byteIndex] &= ~(1 << (7 - bitIndex));
}

static inline HGUInt MyAtoi(const HGChar* c)
{
	HGUInt value = 0;

	const HGChar* p = c;
	while (0 != *p)
	{
		value *= 10;
		value += (HGUInt)(*p - '0');
		p++;
	}

	return value;
}

HGResult HGAPI HGImgFmt_CheckPnmFile(const HGChar* fileName, HGBool* isPnm)
{
	if (NULL == fileName || NULL == isPnm)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	HGPnmLoadInfo info;
	HGResult ret = HGImgFmt_LoadPnmImage(fileName, &info, 0, 0, NULL);
	if (HGBASE_ERR_OK != ret)
	{
		return ret;
	}

	*isPnm = HGTRUE;
	return ret;
}

HGResult HGAPI HGImgFmt_GetPnmTypeFromFileName(const HGChar* fileName, HGUInt* pnmType)
{
	if (NULL == fileName || NULL == pnmType)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	const char* p = strrchr(fileName, '.');
	if (NULL == p)
	{
		return HGBASE_ERR_INVALIDARG;
	}

#if defined(HG_CMP_MSC)
	if (0 == _stricmp(p, ".pbm"))
	{
		*pnmType = HGIMGFMT_PNMTYPE_BINARY_BINARY;
		return HGBASE_ERR_OK;
	}

	if (0 == _stricmp(p, ".pgm"))
	{
		*pnmType = HGIMGFMT_PNMTYPE_GRAY_BINARY;
		return HGBASE_ERR_OK;
	}

	if (0 == _stricmp(p, ".ppm"))
	{
		*pnmType = HGIMGFMT_PNMTYPE_RGB_BINARY;
		return HGBASE_ERR_OK;
	}
#else
	if (0 == strcasecmp(p, ".pbm"))
	{
		*pnmType = HGIMGFMT_PNMTYPE_BINARY_BINARY;
		return HGBASE_ERR_OK;
	}

	if (0 == strcasecmp(p, ".pgm"))
	{
		*pnmType = HGIMGFMT_PNMTYPE_GRAY_BINARY;
		return HGBASE_ERR_OK;
	}

	if (0 == strcasecmp(p, ".ppm"))
	{
		*pnmType = HGIMGFMT_PNMTYPE_RGB_BINARY;
		return HGBASE_ERR_OK;
	}
#endif

	return HGBASE_ERR_FAIL;
}

static HGResult PnmLoadInfo(FILE* file, HGUInt pnmType, HGUInt *width, HGUInt *height, HGUInt *maxColor)
{
	assert(NULL != file && NULL != width && NULL != height && NULL != maxColor);

	bool getWidth = false;
	bool getHeight = false;
	bool getMaxColor = false;
	char buf[256];
	int bufLen = 0;

	while (1)
	{
		HGByte c = 0;
		if (1 != fread(&c, 1, 1, file))
		{
			return HGBASE_ERR_FILEERROR;
		}

		if (c == '#' || c == '\n' || c == '\r' || c == '\t' || c == ' ')
		{
			if (0 != bufLen)
			{
				buf[bufLen] = 0;

				if (!getWidth)
				{
					*width = MyAtoi(buf);
					if (*width == 0)
						return HGBASE_ERR_FILEERROR;
					getWidth = true;
				}
				else if (!getHeight)
				{
					*height = MyAtoi(buf);
					if (*height == 0)
						return HGBASE_ERR_FILEERROR;
					getHeight = true;
				}
				else if (!getMaxColor)
				{
					*maxColor = MyAtoi(buf);
					if (*maxColor == 0)
						return HGBASE_ERR_FILEERROR;
					getMaxColor = true;
				}

				bufLen = 0;
			}

			if (c == '#')
			{
				while (c != '\n')
				{
					if (1 != fread(&c, 1, 1, file))
					{
						return HGBASE_ERR_FILEERROR;
					}
				}
			}
		}
		else if (c >= '0' && c <= '9')
		{
			buf[bufLen] = c;
			++bufLen;
		}
		else
		{
			return HGBASE_ERR_FILEERROR;
		}

		if (pnmType != HGIMGFMT_PNMTYPE_BINARY_ASCII && pnmType != HGIMGFMT_PNMTYPE_BINARY_BINARY)
		{
			if (getMaxColor)
			{
				break;
			}
		}
		else
		{
			if (getHeight)
			{
				break;
			}
		}
	}

	return HGBASE_ERR_OK;
}

static HGResult PnmLoadImage(FILE* file, HGUInt pnmType, HGPnmLoadInfo* info, HGUInt imgType, HGUInt imgOrigin, HGImage* image)
{
	HGUInt width = 0, height = 0, maxColor = 0;
	HGResult ret = PnmLoadInfo(file, pnmType, &width, &height, &maxColor);
	if (HGBASE_ERR_OK != ret)
	{
		return ret;
	}

	if (pnmType != HGIMGFMT_PNMTYPE_BINARY_ASCII && pnmType != HGIMGFMT_PNMTYPE_BINARY_BINARY)
	{
		if (maxColor != 255)
			return HGBASE_ERR_FILEERROR;
	}
	else
	{
		maxColor = 1;
	}

	if (NULL != info)
	{
		info->width = width;
		info->height = height;
		info->type = pnmType;
	}

	if (NULL != image)
	{
		if (imgType == 0)
		{
			if (pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII || pnmType == HGIMGFMT_PNMTYPE_BINARY_BINARY)
				imgType = HGBASE_IMGTYPE_BINARY;
			else if (pnmType == HGIMGFMT_PNMTYPE_GRAY_ASCII || pnmType == HGIMGFMT_PNMTYPE_GRAY_BINARY)
				imgType = HGBASE_IMGTYPE_GRAY;
			else if (pnmType == HGIMGFMT_PNMTYPE_RGB_ASCII || pnmType == HGIMGFMT_PNMTYPE_RGB_BINARY)
				imgType = HGBASE_IMGTYPE_RGB;
		}

		if (imgOrigin == 0)
		{
			imgOrigin = HGBASE_IMGORIGIN_TOP;
		}

		HGImage image2 = NULL;

		if (pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII || pnmType == HGIMGFMT_PNMTYPE_BINARY_BINARY)
			ret = HGBase_CreateImage(width, height, HGBASE_IMGTYPE_BINARY, HGBASE_IMGORIGIN_TOP, &image2);
		else if (pnmType == HGIMGFMT_PNMTYPE_GRAY_ASCII || pnmType == HGIMGFMT_PNMTYPE_GRAY_BINARY)
			ret = HGBase_CreateImage(width, height, HGBASE_IMGTYPE_GRAY, HGBASE_IMGORIGIN_TOP, &image2);
		else if (pnmType == HGIMGFMT_PNMTYPE_RGB_ASCII || pnmType == HGIMGFMT_PNMTYPE_RGB_BINARY)
			ret = HGBase_CreateImage(width, height, HGBASE_IMGTYPE_RGB, HGBASE_IMGORIGIN_TOP, &image2);
		if (HGBASE_ERR_OK != ret)
		{
			return ret;
		}

		uint8_t* data;
		HGBase_GetImageData(image2, &data);
		HGImageInfo imgInfo;
		HGBase_GetImageInfo(image2, &imgInfo);

		if (pnmType == HGIMGFMT_PNMTYPE_BINARY_BINARY || pnmType == HGIMGFMT_PNMTYPE_GRAY_BINARY || pnmType == HGIMGFMT_PNMTYPE_RGB_BINARY)
		{
			HGUInt lineSize = ((width + 7) & ~7) >> 3;
			if (pnmType == HGIMGFMT_PNMTYPE_GRAY_BINARY)
				lineSize = width;
			else if (pnmType == HGIMGFMT_PNMTYPE_RGB_BINARY)
				lineSize = width * 3;

			for (HGUInt i = 0; i < height; ++i)
			{
				if (lineSize != fread(data + i * imgInfo.widthStep, 1, lineSize, file))
				{
					HGBase_DestroyImage(image2);
					return HGBASE_ERR_FILEERROR;
				}
			}

			if (pnmType == HGIMGFMT_PNMTYPE_BINARY_BINARY)
			{
				HGBase_ReverseImage(image2, image2);
			}
		}
		else if (pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII || pnmType == HGIMGFMT_PNMTYPE_GRAY_ASCII || pnmType == HGIMGFMT_PNMTYPE_RGB_ASCII)
		{
			HGChar readBuf[2048];
			HGUInt readBufLen = 0;
			HGChar* curReadPos = readBuf;
			HGChar buf[256];
			HGUInt bufLen = 0;
			HGByte* p = data;
			HGUInt idx = 0;

			while (1)
			{
				if (0 == readBufLen)
				{
					readBufLen = (HGUInt)fread(readBuf, 1, 2048, file);
					curReadPos = readBuf;
				}

				HGByte c = 0;
				HGBool getChar = HGFALSE;
				if (0 != readBufLen)
				{
					c = *curReadPos;
					++curReadPos;
					--readBufLen;
					getChar = HGTRUE;
				}

				if (!getChar || c == '#' || c == '\n' || c == '\r' || c == '\t' || c == ' ')
				{
					if (0 != bufLen)
					{
						buf[bufLen] = 0;
						HGUInt pixel = MyAtoi(buf);
						if (pixel < 0 || pixel > maxColor)
						{
							HGBase_DestroyImage(image2);
							return HGBASE_ERR_FILEERROR;
						}

						if (pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII)
							SetBit(p, idx, (pixel == 0) ? 1 : 0);
						else
							p[idx] = (HGByte)pixel;

						HGUInt step = width;
						if (pnmType == HGIMGFMT_PNMTYPE_RGB_ASCII)
							step = width * 3;

						++idx;
						if (idx == step)
						{
							p += imgInfo.widthStep;
							idx = 0;
						}

						bufLen = 0;
					}

					if (!getChar)
					{
						if (p == data + height * imgInfo.widthStep)
						{
							break;
						}
						else
						{
							HGBase_DestroyImage(image2);
							return HGBASE_ERR_FILEERROR;
						}
					}
					else if (c == '#')
					{
						while (c != '\n')
						{
							c = 0;
							HGBool getChar = HGFALSE;
							if (0 != readBufLen)
							{
								c = *curReadPos;
								++curReadPos;
								--readBufLen;
								getChar = HGTRUE;
							}

							if (!getChar)
							{
								HGBase_DestroyImage(image2);
								return HGBASE_ERR_FILEERROR;
							}
						}
					}
				}
				else if (c >= '0' && c <= '9')
				{
					buf[bufLen] = c;
					++bufLen;
				}
				else
				{
					HGBase_DestroyImage(image2);
					return HGBASE_ERR_FILEERROR;
				}

				if (p == data + height * imgInfo.widthStep)
				{
					break;
				}
			}
		}

		if (imgInfo.type == imgType && imgInfo.origin == imgOrigin)
		{
			*image = image2;
		}
		else
		{
			ret = HGBase_CloneImage(image2, imgType, imgOrigin, image);
			HGBase_DestroyImage(image2);
			if (HGBASE_ERR_OK != ret)
			{
				return ret;
			}
		}
	}

	return HGBASE_ERR_OK;
}

HGResult HGAPI HGImgFmt_LoadPnmImage(const HGChar* fileName, HGPnmLoadInfo* info,
	HGUInt imgType, HGUInt imgOrigin, HGImage* image)
{
	if (NULL == fileName)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	if (NULL == image)
	{
		if (0 != imgType || 0 != imgOrigin)
		{
			return HGBASE_ERR_INVALIDARG;
		}
	}
	else
	{
		if (0 != imgType && HGBASE_IMGTYPE_BINARY != imgType && HGBASE_IMGTYPE_GRAY != imgType
			&& HGBASE_IMGTYPE_BGR != imgType && HGBASE_IMGTYPE_RGB != imgType
			&& HGBASE_IMGTYPE_BGRA != imgType && HGBASE_IMGTYPE_RGBA != imgType)
		{
			return HGBASE_ERR_INVALIDARG;
		}

		if (0 != imgOrigin && HGBASE_IMGORIGIN_TOP != imgOrigin && HGBASE_IMGORIGIN_BOTTOM != imgOrigin)
		{
			return HGBASE_ERR_INVALIDARG;
		}
	}

#if defined(HG_CMP_MSC)
	if (0 != _access(fileName, 0))
#else
	if (0 != access(fileName, 0))
#endif
	{
		return HGBASE_ERR_FILENOTEXIST;
	}

	FILE* file = fopen(fileName, "rb");
	if (NULL == file)
	{
		return HGBASE_ERR_ACCESSDENIED;
	}

	HGByte magicKey[2] = {0};
	if (2 != fread(magicKey, 1, 2, file))
	{
		fclose(file);
		return HGBASE_ERR_FILEERROR;
	}

	HGUInt pnmType = 0;
	if (magicKey[0] == 'P' && (magicKey[1] == '1'))
	{
		pnmType = HGIMGFMT_PNMTYPE_BINARY_ASCII;
	}
	else if (magicKey[0] == 'P' && (magicKey[1] == '2'))
	{
		pnmType = HGIMGFMT_PNMTYPE_GRAY_ASCII;
	}
	else if (magicKey[0] == 'P' && (magicKey[1] == '3'))
	{
		pnmType = HGIMGFMT_PNMTYPE_RGB_ASCII;
	}
	else if (magicKey[0] == 'P' && (magicKey[1] == '4'))
	{
		pnmType = HGIMGFMT_PNMTYPE_BINARY_BINARY;
	}
	else if (magicKey[0] == 'P' && (magicKey[1] == '5'))
	{
		pnmType = HGIMGFMT_PNMTYPE_GRAY_BINARY;
	}
	else if (magicKey[0] == 'P' && (magicKey[1] == '6'))
	{
		pnmType = HGIMGFMT_PNMTYPE_RGB_BINARY;
	}
	
	if (0 == pnmType)
	{
		fclose(file);
		return HGBASE_ERR_FILEERROR;
	}

	HGResult ret = PnmLoadImage(file, pnmType, info, imgType, imgOrigin, image);
	fclose(file);
	return ret;
}

static HGResult PnmSaveImage(HGImage image, const HGChar* fileName, HGUInt pnmType)
{
	FILE* file = fopen(fileName, "wb");
	if (NULL == file)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_ERROR, "PnmSaveImage: fopen fail, %s errno=%d", fileName, errno);
		return HGBASE_ERR_ACCESSDENIED;
	}

	HGByte* data = NULL;
	HGBase_GetImageData(image, &data);
	HGImageInfo imgInfo;
	HGBase_GetImageInfo(image, &imgInfo);

	char magicKey[4] = {0};
	if (HGIMGFMT_PNMTYPE_BINARY_ASCII == pnmType)
	{
		strcpy(magicKey, "P1\n");
	}
	else if (HGIMGFMT_PNMTYPE_GRAY_ASCII == pnmType)
	{
		strcpy(magicKey, "P2\n");
	}
	else if (HGIMGFMT_PNMTYPE_RGB_ASCII == pnmType)
	{
		strcpy(magicKey, "P3\n");
	}
	else if (HGIMGFMT_PNMTYPE_BINARY_BINARY == pnmType)
	{
		strcpy(magicKey, "P4\n");
	}
	else if (HGIMGFMT_PNMTYPE_GRAY_BINARY == pnmType)
	{
		strcpy(magicKey, "P5\n");
	}
	else if (HGIMGFMT_PNMTYPE_RGB_BINARY == pnmType)
	{
		strcpy(magicKey, "P6\n");
	}

	fwrite(magicKey, 1, strlen(magicKey), file);

	char width[20], height[20];
	sprintf(width, "%u\n", imgInfo.width);
	sprintf(height, "%u\n", imgInfo.height);
	fwrite(width, 1, strlen(width), file);
	fwrite(height, 1, strlen(height), file);

	if (HGIMGFMT_PNMTYPE_BINARY_ASCII != pnmType && HGIMGFMT_PNMTYPE_BINARY_BINARY != pnmType)
	{
		char maxColor[] = "255\n";
		fwrite(maxColor, 1, strlen(maxColor), file);
	}

	HGByte* p = data;
	HGInt step = (HGInt)imgInfo.widthStep;
	if (HGBASE_IMGORIGIN_BOTTOM == imgInfo.origin)
	{
		p = data + (imgInfo.height - 1) * imgInfo.widthStep;
		step = -(HGInt)imgInfo.widthStep;
	}

	if (pnmType == HGIMGFMT_PNMTYPE_BINARY_BINARY)
	{
		HGUInt lineSize = ((imgInfo.width + 7) & ~7) >> 3;
		HGByte* buf = (HGByte*)malloc(lineSize);
		if (NULL == buf)
		{
			fclose(file);
			return HGBASE_ERR_OUTOFMEMORY;
		}

		for (HGUInt i = 0; i < imgInfo.height; ++i)
		{
			uint8_t* pEx = p + (HGSize)i * (HGSize)step;
			for (HGUInt j = 0; j < lineSize; ++j)
			{
				buf[j] = ~pEx[j];	// 黑白反色
			}

			fwrite(buf, 1, lineSize, file);
		}

		free(buf);
	}
	else if (pnmType == HGIMGFMT_PNMTYPE_GRAY_BINARY)
	{
		for (HGUInt i = 0; i < imgInfo.height; ++i)
		{
			uint8_t* pEx = p + (HGSize)i * (HGSize)step;
			fwrite(pEx, 1, imgInfo.width, file);
		}
	}
	else if (pnmType == HGIMGFMT_PNMTYPE_RGB_BINARY)
	{
		for (HGUInt i = 0; i < imgInfo.height; ++i)
		{
			uint8_t* pEx = p + (HGSize)i * (HGSize)step;
			fwrite(pEx, 1, imgInfo.width * 3, file);
		}
	}
	else if (pnmType == HGIMGFMT_PNMTYPE_RGB_ASCII || pnmType == HGIMGFMT_PNMTYPE_GRAY_ASCII || pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII)
	{
		HGUInt w = imgInfo.width;
		if (HGIMGFMT_PNMTYPE_RGB_ASCII == pnmType)
			w = imgInfo.width * 3;

		char *buf = (char *)malloc(w * 5 + 1);
		if (NULL == buf)
		{
			fclose(file);
			return HGBASE_ERR_OUTOFMEMORY;
		}

		std::string pixelStr[256];
		for (int i = 0; i < 256; ++i)
		{
			char str[6];
			sprintf(str, "%d", i);
			pixelStr[i] = str;
		}

		for (HGUInt i = 0; i < imgInfo.height; ++i)
		{
			int bufLen = 0;
			int lineSize = 0;

			uint8_t* pEx = p + (HGSize)i * (HGSize)step;
			for (HGUInt j = 0; j < w; ++j)
			{
				HGUInt idx;
				if (pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII)
					idx = (0 == GetBit(pEx, j) ? 1 : 0);
				else
					idx = pEx[j];

				if (0 != lineSize)
				{
					if (lineSize + pixelStr[idx].size() + 1 > 70)
					{
						buf[bufLen] = '\n';
						++bufLen;
						lineSize = 0;
					}
					else
					{
						buf[bufLen] = ' ';
						++bufLen;
						++lineSize;
					}
				}

				strcpy(buf + bufLen, pixelStr[idx].c_str());
				bufLen += (int)pixelStr[idx].size();
				lineSize += (int)pixelStr[idx].size();
			}

			buf[bufLen] = '\n';
			++bufLen;
			lineSize = 0;

			fwrite(buf, 1, bufLen, file);
		}

		free(buf);
	}

	fclose(file);
	return HGBASE_ERR_OK;
}

HGResult HGAPI HGImgFmt_SavePnmImage(HGImage image, const HGPnmSaveInfo* info, const HGChar* fileName)
{
	if (NULL == image || NULL == fileName)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	if (NULL != info)
	{
		if (info->type < HGIMGFMT_PNMTYPE_BINARY_ASCII || info->type > HGIMGFMT_PNMTYPE_RGB_BINARY)
		{
			return HGBASE_ERR_INVALIDARG;
		}
	}

	HGImageInfo imgInfo;
	HGBase_GetImageInfo(image, &imgInfo);

	HGUInt pnmType = HGIMGFMT_PNMTYPE_RGB_BINARY;
	if (HGBASE_IMGTYPE_GRAY == imgInfo.type)
		pnmType = HGIMGFMT_PNMTYPE_GRAY_BINARY;
	else if (HGBASE_IMGTYPE_BINARY == imgInfo.type)
		pnmType = HGIMGFMT_PNMTYPE_BINARY_BINARY;

	if (NULL != info)
	{
		pnmType = info->type;
	}
	else
	{
		HGUInt pnmType2 = 0;
		HGImgFmt_GetPnmTypeFromFileName(fileName, &pnmType2);
		if (0 != pnmType2)
		{
			pnmType = pnmType2;
		}
	}

	HGUInt imgType = HGBASE_IMGTYPE_RGB;
	if (pnmType == HGIMGFMT_PNMTYPE_BINARY_ASCII || pnmType == HGIMGFMT_PNMTYPE_BINARY_BINARY)
		imgType = HGBASE_IMGTYPE_BINARY;
	else if (pnmType == HGIMGFMT_PNMTYPE_GRAY_ASCII || pnmType == HGIMGFMT_PNMTYPE_GRAY_BINARY)
		imgType = HGBASE_IMGTYPE_GRAY;

	if (imgInfo.type != imgType)
	{
		HGImage image2 = NULL;
		HGResult ret = HGBase_CloneImage(image, imgType, HGBASE_IMGORIGIN_TOP, &image2);
		if (ret != HGBASE_ERR_OK)
		{
			return ret;
		}

		ret = PnmSaveImage(image2, fileName, pnmType);
		HGBase_DestroyImage(image2);
		return ret;
	}

	return PnmSaveImage(image, fileName, pnmType);
}