#include "HGOCRTesseract.hpp"
#include "HGOCR.h"
#include "HGOCRRetImpl.hpp"
#include "HGOCRRetImpl.hpp"
#include "../base/HGUtility.h"
#include "../base/HGInfo.h"
#include "../imgfmt/HGBmp.h"
#include "../imgfmt/HGJpeg.h"
#include "HGString.h"

HGOCRTesseract::HGOCRTesseract()
{
	m_baseApi = NULL;
	m_tiffFileName.clear();
	m_tiffWriter = NULL;
}

HGOCRTesseract::~HGOCRTesseract()
{

}

HGResult HGOCRTesseract::Init()
{
	if (NULL != m_baseApi)
	{
		return HGBASE_ERR_FAIL;
	}

	m_baseApi = TessBaseAPICreate();
	if (NULL == m_baseApi)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_ERROR, "HGOCRTesseract::Init: TessBaseAPICreate fail");
		return HGIMGPROC_ERR_OCRINIT;
	}
	
	HGChar moduleName[256];
	HGBase_GetModuleName((void*)HGImgProc_CreateOCRMgr, moduleName, 256);
	HGChar dataPath[256];
	HGBase_GetFilePath(moduleName, dataPath, 256);
	strcat(dataPath, "tessdata");

	int rc = TessBaseAPIInit3(m_baseApi, dataPath, "chi_sim");
	if (0 != rc)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_ERROR, "HGOCRTesseract::Init: TessBaseAPIInit3 fail");
		TessBaseAPIDelete(m_baseApi);
		m_baseApi = NULL;
		return HGIMGPROC_ERR_OCRINIT;
	}

	TessBaseAPISetPageSegMode(m_baseApi, TessPageSegMode::PSM_AUTO_OSD);
	return HGBASE_ERR_OK;
}

HGResult HGOCRTesseract::Deinit()
{
	if (NULL == m_baseApi)
	{
		return HGBASE_ERR_FAIL;
	}

	ClearImageList();

	TessBaseAPIDelete(m_baseApi);
	m_baseApi = NULL;
	return HGBASE_ERR_OK;
}

HGResult HGOCRTesseract::ImageOCR(HGImage image, class HGOCRRetImpl** ocrRet)
{
	if (NULL == image || NULL == ocrRet)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	HGImage image2 = NULL;

	HGImageInfo imgInfo;
	HGBase_GetImageInfo(image, &imgInfo);
	if (HGBASE_IMGTYPE_RGB != imgInfo.type || HGBASE_IMGORIGIN_TOP != imgInfo.origin)
	{
		HGResult ret = HGBase_CloneImage(image, HGBASE_IMGTYPE_RGB, HGBASE_IMGORIGIN_TOP, &image2);
		if (HGBASE_ERR_OK != ret)
			return ret;
	}
	else
	{
		image2 = image;
	}

	HGBase_GetImageInfo(image2, &imgInfo);
	HGByte* imageData = NULL;
	HGBase_GetImageData(image2, &imageData);

	TessBaseAPISetImage(m_baseApi, imageData, imgInfo.width, imgInfo.height, 3, imgInfo.widthStep);
	HGUInt xDpi, yDpi;
	HGBase_GetImageDpi(image2, &xDpi, &yDpi);
	TessBaseAPISetSourceResolution(m_baseApi, (xDpi + yDpi) / 2);

	char *text = TessBaseAPIGetUTF8Text(m_baseApi);
	if (NULL == text)
	{
		if (image2 != image)
			HGBase_DestroyImage(image2);
		return HGIMGPROC_ERR_OCR;
	}

	std::vector<std::string> blockInfo;
	blockInfo.push_back(Utf8ToStdString(text));
	*ocrRet = new HGOCRRetImpl(blockInfo);

	TessDeleteText(text);
	if (image2 != image)
		HGBase_DestroyImage(image2);
	return HGBASE_ERR_OK;
}

HGResult HGOCRTesseract::ImageOCRToFile(HGImage image, HGUInt outType, const HGChar* outFileName)
{
	if (NULL == image || outType > HGIMGPROC_OCROUTTYPE_OFD || NULL == outFileName)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	if (0 == outType)
		outType = GetOutTypeByFileName(outFileName);

	HGChar tmpFileName[256];
	HGBase_GetTmpFileName("bmp", tmpFileName, 256);
	HGResult ret = HGImgFmt_SaveBmpImage(image, NULL, tmpFileName);
	if (HGBASE_ERR_OK != ret)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_ERROR, "HGOCRTesseract::ImageOCRToFile: HGImgFmt_SaveBmpImage fail %s", tmpFileName);
		return ret;
	}

	ret = OCRToFile(tmpFileName, outType, outFileName);
	HGBase_DeleteFile(tmpFileName);
	return ret;
}

HGResult HGOCRTesseract::ImageTextDirectOCR(HGImage image, HGUInt* direct)
{
	if (NULL == image || NULL == direct)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	HGImage image2 = NULL;

	HGImageInfo imgInfo;
	HGBase_GetImageInfo(image, &imgInfo);
	if (HGBASE_IMGTYPE_RGB != imgInfo.type || HGBASE_IMGORIGIN_TOP != imgInfo.origin)
	{
		HGResult ret = HGBase_CloneImage(image, HGBASE_IMGTYPE_RGB, HGBASE_IMGORIGIN_TOP, &image2);
		if (HGBASE_ERR_OK != ret)
			return ret;
	}
	else
	{
		image2 = image;
	}

	HGBase_GetImageInfo(image2, &imgInfo);
	HGByte* imageData = NULL;
	HGBase_GetImageData(image2, &imageData);

	TessBaseAPISetImage(m_baseApi, imageData, imgInfo.width, imgInfo.height, 3, imgInfo.widthStep);
	HGUInt xDpi, yDpi;
	HGBase_GetImageDpi(image2, &xDpi, &yDpi);
	TessBaseAPISetSourceResolution(m_baseApi, (xDpi + yDpi) / 2);

	int orientation = MyOSD(m_baseApi);
	if (TessOrientation::ORIENTATION_PAGE_UP == orientation)
		*direct = HGIMGPROC_OCRTEXTDIRECT_ORI;
	else if (TessOrientation::ORIENTATION_PAGE_RIGHT == orientation)
		*direct = HGIMGPROC_OCRTEXTDIRECT_RIGHT;
	else if (TessOrientation::ORIENTATION_PAGE_DOWN == orientation)
		*direct = HGIMGPROC_OCRTEXTDIRECT_180;
	else if (TessOrientation::ORIENTATION_PAGE_LEFT == orientation)
		*direct = HGIMGPROC_OCRTEXTDIRECT_LEFT;

	if (image2 != image)
		HGBase_DestroyImage(image2);
	return HGBASE_ERR_OK;
}

HGResult HGOCRTesseract::AddToImageList(HGImage image)
{
	if (NULL == image)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	if (NULL == m_tiffWriter)
	{
		HGChar tmpFileName[256] = { 0 };
		HGBase_GetTmpFileName("tif", tmpFileName, 256);
		m_tiffFileName = tmpFileName;

		HGResult ret = HGImgFmt_OpenTiffWriter(m_tiffFileName.c_str(), &m_tiffWriter);
		if (HGBASE_ERR_OK != ret)
		{
			HGBase_WriteInfo(HGBASE_INFOTYPE_ERROR, "HGOCRTesseract::AddToImageList: HGImgFmt_OpenTiffWriter fail %s", m_tiffFileName.c_str());
			m_tiffFileName.clear();
			return ret;
		}
	}

	return HGImgFmt_SaveImageToTiffWriter(m_tiffWriter, image, NULL);
}

HGResult HGOCRTesseract::ClearImageList()
{
	if (NULL != m_tiffWriter)
	{
		HGImgFmt_CloseTiffWriter(m_tiffWriter);
		m_tiffWriter = NULL;
		HGBase_DeleteFile(m_tiffFileName.c_str());
		m_tiffFileName.clear();
	}

	return HGBASE_ERR_OK;
}

HGResult HGOCRTesseract::ImageListOCRToFile(HGUInt outType, const HGChar* outFileName, HGImageListOcrFunc func, HGPointer param)
{
	if (NULL == m_tiffWriter)
	{
		return HGBASE_ERR_FAIL;
	}

	if (outType > HGIMGPROC_OCROUTTYPE_OFD || NULL == outFileName)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	if (0 == outType)
		outType = GetOutTypeByFileName(outFileName);

	if (HGIMGPROC_OCROUTTYPE_PDF != outType && HGIMGPROC_OCROUTTYPE_TXT != outType)
	{
		return HGBASE_ERR_INVALIDARG;
	}

	HGImgFmt_CloseTiffWriter(m_tiffWriter);
	m_tiffWriter = NULL;
	assert(!m_tiffFileName.empty());
	HGResult ret = OCRToFile(m_tiffFileName.c_str(), outType, outFileName);
	HGBase_DeleteFile(m_tiffFileName.c_str());
	m_tiffFileName.clear();
	return ret;
}

HGResult HGOCRTesseract::OCRToFile(const HGChar* inFileName, HGUInt outType, const HGChar* outFileName)
{
	assert(NULL != inFileName);
	assert(NULL != outFileName);

	HGResult ret = HGBASE_ERR_NOTSUPPORT;
	if (HGIMGPROC_OCROUTTYPE_PDF == outType)
	{
		HGChar outputbase[256] = { 0 };
		const char* p = strrchr(outFileName, '.');
		if (NULL != p && 0 == strcmp(p, ".pdf"))
			memcpy(outputbase, outFileName, p - outFileName);
		else
			strcpy(outputbase, outFileName);

		ret = HGIMGPROC_ERR_OCR;
		TessResultRenderer* pdfRender = TessPDFRendererCreate(outputbase, TessBaseAPIGetDatapath(m_baseApi), FALSE);
		if (NULL != pdfRender)
		{
			if (TessBaseAPIProcessPages(m_baseApi, inFileName, NULL, 0, pdfRender))
				ret = HGBASE_ERR_OK;
			TessDeleteResultRenderer(pdfRender);

			if (HGBASE_ERR_OK == ret)
			{
				HGChar destFileName[256];
				sprintf(destFileName, "%s.pdf", outputbase);
#if defined(HG_CMP_MSC)
				MoveFileA(destFileName, outFileName);
#else
				rename(destFileName, outFileName);
#endif
			}
		}
	}
	else if (HGIMGPROC_OCROUTTYPE_TXT == outType)
	{
		HGChar outputbase[256] = { 0 };
		const char* p = strrchr(outFileName, '.');
		if (NULL != p && 0 == strcmp(p, ".txt"))
			memcpy(outputbase, outFileName, p - outFileName);
		else
			strcpy(outputbase, outFileName);

		ret = HGIMGPROC_ERR_OCR;
		TessResultRenderer* txtRender = TessTextRendererCreate(outputbase);
		if (NULL != txtRender)
		{
			if (TessBaseAPIProcessPages(m_baseApi, inFileName, NULL, 0, txtRender))
				ret = HGBASE_ERR_OK;
			TessDeleteResultRenderer(txtRender);

			if (HGBASE_ERR_OK == ret)
			{
				HGChar destFileName[256];
				sprintf(destFileName, "%s.txt", outputbase);
#if defined(HG_CMP_MSC)
				MoveFileA(destFileName, outFileName);
#else
				rename(destFileName, outFileName);
#endif
			}
		}
	}

	return ret;
}