#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_dll = NULL;
	m_baseApi = NULL;
	m_tiffFileName.clear();
	m_tiffWriter = NULL;
}

HGOCRTesseract::~HGOCRTesseract()
{

}

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

	assert(NULL == m_dll);

	HGChar moduleName[256];
	HGBase_GetModuleName((void*)HGImgProc_CreateOCRMgr, moduleName, 256);
	HGChar dllPath[256];
	HGBase_GetFilePath(moduleName, dllPath, 256);

#if defined(HG_CMP_MSC)
	strcat(dllPath, "tesseract-ocr\\libtesseract-5.dll");
#else
	strcat(dllPath, "tesseract-ocr\\libtesseract-5.so");
#endif
	HGResult ret = HGBase_CreateDll(dllPath, &m_dll);
	if (HGBASE_ERR_OK != ret)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_DESC, "Load OCR Library Failed, path=%s", dllPath);
		return ret;
	}

	m_baseApi = TessBaseAPICreate();
	if (NULL == m_baseApi)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_DESC, "Init OCR Library Failed");
		HGBase_DestroyDll(m_dll);
		m_dll = NULL;
		return HGIMGPROC_ERR_FAIL;
	}
	
	HGChar dataPath[256];
	HGBase_GetFilePath(moduleName, dataPath, 256);
	strcat(dataPath, "tesseract-ocr\\tessdata");

	int rc = TessBaseAPIInit3(m_baseApi, dataPath, "chi_sim");
	if (0 != rc)
	{
		HGBase_WriteInfo(HGBASE_INFOTYPE_DESC, "Init OCR Library Failed 2");
		TessBaseAPIDelete(m_baseApi);
		m_baseApi = NULL;
		HGBase_DestroyDll(m_dll);
		m_dll = NULL;
		return HGIMGPROC_ERR_FAIL;
	}

	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;
	HGBase_DestroyDll(m_dll);
	m_dll = 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);

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

	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);
	HGImgFmt_SaveBmpImage(image, NULL, tmpFileName);

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

	TessPageIterator* iter = TessBaseAPIAnalyseLayout(m_baseApi);
	if (NULL == iter)
	{
		if (image2 != image)
			HGBase_DestroyImage(image2);
		return HGIMGPROC_ERR_FAIL;
	}

	TessOrientation orientation;
	TessWritingDirection writing_direction;
	TessTextlineOrder textline_order;
	float deskew_angle;
	TessPageIteratorOrientation(iter, &orientation, &writing_direction, &textline_order, &deskew_angle);

	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;

	TessPageIteratorDelete(iter);
	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)
		{
			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 HGIMGPROC_ERR_FAIL;
	}

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

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

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

TessBaseAPI* HGOCRTesseract::TessBaseAPICreate()
{
	typedef TessBaseAPI* (*Func)();
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPICreate", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func();
}

void HGOCRTesseract::TessBaseAPIDelete(TessBaseAPI* handle)
{
	typedef void (*Func)(TessBaseAPI *);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIDelete", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle);
}

void HGOCRTesseract::TessBaseAPISetInputName(TessBaseAPI* handle, const char* name)
{
	typedef void (*Func)(TessBaseAPI*, const char *);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPISetInputName", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle, name);
}

void HGOCRTesseract::TessBaseAPISetOutputName(TessBaseAPI* handle, const char* name)
{
	typedef void (*Func)(TessBaseAPI*, const char*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPISetOutputName", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle, name);
}

char* HGOCRTesseract::TessBaseAPIGetUTF8Text(TessBaseAPI* handle)
{
	typedef char* (*Func)(TessBaseAPI*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIGetUTF8Text", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func(handle);
}

int HGOCRTesseract::TessBaseAPIInit3(TessBaseAPI* handle, const char* datapath, const char* language)
{
	typedef int (*Func)(TessBaseAPI*, const char *, const char *);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIInit3", (HGPointer*)&func);
	if (NULL == func)
	{
		return -1;
	}

	return func(handle, datapath, language);
}

BOOL HGOCRTesseract::TessBaseAPIProcessPages(TessBaseAPI* handle, const char* filename, const char* retry_config,
	int timeout_millisec, TessResultRenderer* renderer)
{
	typedef BOOL (*Func)(TessBaseAPI*, const char*, const char*, int , TessResultRenderer*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIProcessPages", (HGPointer*)&func);
	if (NULL == func)
	{
		return FALSE;
	}

	return func(handle, filename, retry_config, timeout_millisec, renderer);
}

void HGOCRTesseract::TessDeleteText(const char* text)
{
	typedef void (*Func)(const char *);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessDeleteText", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(text);
}

void HGOCRTesseract::TessBaseAPISetPageSegMode(TessBaseAPI* handle, TessPageSegMode mode)
{
	typedef void (*Func)(TessBaseAPI*, TessPageSegMode);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPISetPageSegMode", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle, mode);
}

BOOL HGOCRTesseract::TessBaseAPISetVariable(TessBaseAPI* handle, const char* name, const char* value)
{
	typedef BOOL (*Func)(TessBaseAPI*, const char*, const char*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPISetVariable", (HGPointer*)&func);
	if (NULL == func)
	{
		return FALSE;
	}

	return func(handle, name, value);
}

TessPageIterator* HGOCRTesseract::TessBaseAPIAnalyseLayout(TessBaseAPI* handle)
{
	typedef TessPageIterator* (*Func)(TessBaseAPI*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIAnalyseLayout", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func(handle);
}

void HGOCRTesseract::TessPageIteratorDelete(TessPageIterator* handle)
{
	typedef void (*Func)(TessPageIterator*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessPageIteratorDelete", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle);
}

void HGOCRTesseract::TessPageIteratorOrientation(TessPageIterator* handle, TessOrientation* orientation, TessWritingDirection* writing_direction,
	TessTextlineOrder* textline_order, float* deskew_angle)
{
	typedef void (*Func)(TessPageIterator*, TessOrientation*, TessWritingDirection*, TessTextlineOrder*, float*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessPageIteratorOrientation", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle, orientation, writing_direction, textline_order, deskew_angle);
}

void HGOCRTesseract::TessBaseAPISetImage(TessBaseAPI* handle, const unsigned char* imagedata, int width, int height,
	int bytes_per_pixel, int bytes_per_line)
{
	typedef void (*Func)(TessBaseAPI*, const unsigned char*, int, int, int, int);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPISetImage", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(handle, imagedata, width, height, bytes_per_pixel, bytes_per_line);
}

int HGOCRTesseract::TessBaseAPIRecognize(TessBaseAPI* handle, ETEXT_DESC* monitor)
{
	typedef int (*Func)(TessBaseAPI*, ETEXT_DESC*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIRecognize", (HGPointer*)&func);
	if (NULL == func)
	{
		return -1;
	}

	return func(handle, monitor);
}

const char* HGOCRTesseract::TessBaseAPIGetDatapath(TessBaseAPI* handle)
{
	typedef const char* (*Func)(TessBaseAPI*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessBaseAPIGetDatapath", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func(handle);
}

TessResultRenderer* HGOCRTesseract::TessPDFRendererCreate(const char* outputbase, const char* datadir, BOOL textonly)
{
	typedef TessResultRenderer* (*Func)(const char*, const char*, BOOL);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessPDFRendererCreate", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func(outputbase, datadir, textonly);
}

void HGOCRTesseract::TessDeleteResultRenderer(TessResultRenderer* renderer)
{
	typedef void (*Func)(TessResultRenderer*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessDeleteResultRenderer", (HGPointer*)&func);
	if (NULL == func)
	{
		return;
	}

	return func(renderer);
}

TessResultRenderer* HGOCRTesseract::TessTextRendererCreate(const char* outputbase)
{
	typedef TessResultRenderer* (*Func)(const char*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessTextRendererCreate", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func(outputbase);
}

const char* HGOCRTesseract::TessResultRendererExtention(TessResultRenderer* renderer)
{
	typedef const char* (*Func)(TessResultRenderer*);
	Func func = NULL;
	HGBase_GetDllProcAddress(m_dll, "TessResultRendererExtention", (HGPointer*)&func);
	if (NULL == func)
	{
		return NULL;
	}

	return func(renderer);
}