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

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);
	*isPnm = (HGBASE_ERR_OK == ret ? HGTRUE : HGFALSE);
	return HGBASE_ERR_OK;
}

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

	bool getWidth = false;
	bool getHeight = false;
	bool getMaxColor = false;
	std::vector<char> buf;

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

		if (c == '#' || c == '\n' || c == '\r' || c == '\t' || c == ' ')
		{
			if (!buf.empty())
			{
				buf.push_back(0);

				if (!getWidth)
				{
					*width = atoi(&buf[0]);
					if (*width == 0)
						return HGBASE_ERR_INVALIDDATA;
					getWidth = true;
				}
				else if (!getHeight)
				{
					*height = atoi(&buf[0]);
					if (*height == 0)
						return HGBASE_ERR_INVALIDDATA;
					getHeight = true;
				}
				else if (!getMaxColor)
				{
					*maxColor = atoi(&buf[0]);
					if (*maxColor == 0)
						return HGBASE_ERR_INVALIDDATA;
					getMaxColor = true;
				}

				buf.clear();
			}

			if (c == '#')
			{
				while (c != '\n')
				{
					if (1 != fread(&c, 1, 1, file))
					{
						return HGBASE_ERR_INVALIDDATA;
					}
				}
			}
		}
		else if (c >= '0' && c <= '9')
		{
			buf.push_back(c);
		}
		else
		{
			return HGBASE_ERR_INVALIDDATA;
		}

		if (getMaxColor)
		{
			break;
		}
	}

	return HGBASE_ERR_OK;
}

static HGResult BnmLoadP6Image(FILE* file, HGPnmLoadInfo* info, HGUInt imgType, HGUInt imgOrigin, HGImage* image)
{
	HGUInt width, height, maxColor;
	HGResult ret = BnmLoadP6Info(file , &width, &height, &maxColor);
	if (HGBASE_ERR_OK != ret || 255 != maxColor)
	{
		return ret;
	}

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

	if (NULL != image)
	{
		if (imgType == 0)
		{
			imgType = HGBASE_IMGTYPE_RGB;
		}

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

		HGImage image2 = NULL;
		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);

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

		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;
		}
	}

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

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

	HGResult ret = HGBASE_ERR_INVALIDDATA;
	if (magicKey[0] == 'P' && magicKey[1] == '6')
	{
		ret = BnmLoadP6Image(file, info, imgType, imgOrigin, image);
	}

	fclose(file);
	return ret;
}

static HGResult BnmSaveP6Image(HGImage image, const HGChar* fileName)
{
	HGImageInfo imgInfo;
	HGBase_GetImageInfo(image, &imgInfo);
	assert(imgInfo.type == HGBASE_IMGTYPE_RGB);

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

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

	char 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);

	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;
	}

	for (HGUInt i = 0; i < imgInfo.height; ++i)
	{
		uint8_t* pEx = p + (HGSize)i * (HGSize)step;
		fwrite(pEx, 1, imgInfo.width * 3, file);
	}

	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_RGB_BINRAY)
		{
			return HGBASE_ERR_INVALIDARG;
		}
	}

	HGImageInfo imgInfo;
	HGBase_GetImageInfo(image, &imgInfo);
	if (imgInfo.type != HGBASE_IMGTYPE_RGB)
	{
		return HGBASE_ERR_INVALIDDATA;
	}

	return BnmSaveP6Image(image, fileName);
}