# 调用 PaddleOCR-json.exe 的 Python Api
# 项目主页：
# https://github.com/hiroi-sora/PaddleOCR-json

import os
import socket  # 套接字
import atexit  # 退出处理
import subprocess  # 进程，管道
import re  # regex
from json import loads as jsonLoads, dumps as jsonDumps
from sys import platform as sysPlatform  # popen静默模式
from base64 import b64encode  # base64 编码


class PPOCR_pipe:  # 调用OCR（管道模式）
    def __init__(self, exePath: str, modelsPath: str = None, argument: dict = None):
        """初始化识别器（管道模式）。\n
        `exePath`: 识别器`PaddleOCR_json.exe`的路径。\n
        `modelsPath`: 识别库`models`文件夹的路径。若为None则默认识别库与识别器在同一目录下。\n
        `argument`: 启动参数，字典`{"键":值}`。参数说明见 https://github.com/hiroi-sora/PaddleOCR-json
        """
        # 私有成员变量
        self.__ENABLE_CLIPBOARD = False

        exePath = os.path.abspath(exePath)
        cwd = os.path.abspath(os.path.join(exePath, os.pardir))  # 获取exe父文件夹
        cmds = [exePath]
        # 处理启动参数
        if modelsPath is not None:
            if os.path.exists(modelsPath) and os.path.isdir(modelsPath):
                cmds += ["--models_path", os.path.abspath(modelsPath)]
            else:
                raise Exception(
                    f"Input modelsPath doesn't exits or isn't a directory. modelsPath: [{modelsPath}]"
                )
        if isinstance(argument, dict):
            for key, value in argument.items():
                # Popen() 要求输入list里所有的元素都是 str 或 bytes
                if isinstance(value, bool):
                    cmds += [f"--{key}={value}"]  # 布尔参数必须键和值连在一起
                elif isinstance(value, str):
                    cmds += [f"--{key}", value]
                else:
                    cmds += [f"--{key}", str(value)]
        # 设置子进程启用静默模式，不显示控制台窗口
        self.ret = None
        startupinfo = None
        if "win32" in str(sysPlatform).lower():
            startupinfo = subprocess.STARTUPINFO()
            startupinfo.dwFlags = (
                subprocess.CREATE_NEW_CONSOLE | subprocess.STARTF_USESHOWWINDOW
            )
            startupinfo.wShowWindow = subprocess.SW_HIDE
        self.ret = subprocess.Popen(  # 打开管道
            cmds,
            cwd=cwd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.DEVNULL,  # 丢弃stderr的内容
            startupinfo=startupinfo,  # 开启静默模式
        )
        # 启动子进程
        while True:
            if not self.ret.poll() == None:  # 子进程已退出，初始化失败
                raise Exception(f"OCR init fail.")
            initStr = self.ret.stdout.readline().decode("utf-8", errors="ignore")
            if "OCR init completed." in initStr:  # 初始化成功
                break
            elif "OCR clipboard enbaled." in initStr:  # 检测到剪贴板已启用
                self.__ENABLE_CLIPBOARD = True
        atexit.register(self.exit)  # 注册程序终止时执行强制停止子进程

    def isClipboardEnabled(self) -> bool:
        return self.__ENABLE_CLIPBOARD

    def getRunningMode(self) -> str:
        # 默认管道模式只能运行在本地
        return "local"

    def runDict(self, writeDict: dict):
        """传入指令字典，发送给引擎进程。\n
        `writeDict`: 指令字典。\n
        `return`:  {"code": 识别码, "data": 内容列表或错误信息字符串}\n"""
        # 检查子进程
        if not self.ret:
            return {"code": 901, "data": f"引擎实例不存在。"}
        if not self.ret.poll() == None:
            return {"code": 902, "data": f"子进程已崩溃。"}
        # 输入信息
        writeStr = jsonDumps(writeDict, ensure_ascii=True, indent=None) + "\n"
        try:
            self.ret.stdin.write(writeStr.encode("utf-8"))
            self.ret.stdin.flush()
        except Exception as e:
            return {
                "code": 902,
                "data": f"向识别器进程传入指令失败，疑似子进程已崩溃。{e}",
            }
        # 获取返回值
        try:
            getStr = self.ret.stdout.readline().decode("utf-8", errors="ignore")
        except Exception as e:
            return {"code": 903, "data": f"读取识别器进程输出值失败。异常信息：[{e}]"}
        try:
            return jsonLoads(getStr)
        except Exception as e:
            return {
                "code": 904,
                "data": f"识别器输出值反序列化JSON失败。异常信息：[{e}]。原始内容：[{getStr}]",
            }

    def run(self, imgPath: str):
        """对一张本地图片进行文字识别。\n
        `exePath`: 图片路径。\n
        `return`:  {"code": 识别码, "data": 内容列表或错误信息字符串}\n"""
        writeDict = {"image_path": imgPath}
        return self.runDict(writeDict)

    def runClipboard(self):
        """立刻对剪贴板第一位的图片进行文字识别。\n
        `return`:  {"code": 识别码, "data": 内容列表或错误信息字符串}\n"""
        if self.__ENABLE_CLIPBOARD:
            return self.run("clipboard")
        else:
            raise Exception("剪贴板功能不存在或已禁用。")

    def runBase64(self, imageBase64: str):
        """对一张编码为base64字符串的图片进行文字识别。\n
        `imageBase64`: 图片base64字符串。\n
        `return`:  {"code": 识别码, "data": 内容列表或错误信息字符串}\n"""
        writeDict = {"image_base64": imageBase64}
        return self.runDict(writeDict)

    def runBytes(self, imageBytes):
        """对一张图片的字节流信息进行文字识别。\n
        `imageBytes`: 图片字节流。\n
        `return`:  {"code": 识别码, "data": 内容列表或错误信息字符串}\n"""
        imageBase64 = b64encode(imageBytes).decode("utf-8")
        return self.runBase64(imageBase64)

    def exit(self):
        """关闭引擎子进程"""
        if hasattr(self, "ret"):
            if not self.ret:
                return
            try:
                self.ret.kill()  # 关闭子进程
            except Exception as e:
                print(f"[Error] ret.kill() {e}")
        self.ret = None
        atexit.unregister(self.exit)  # 移除退出处理
        print("###  PPOCR引擎子进程关闭！")

    @staticmethod
    def printResult(res: dict):
        """用于调试，格式化打印识别结果。\n
        `res`: OCR识别结果。"""

        # 识别成功
        if res["code"] == 100:
            index = 1
            for line in res["data"]:
                print(
                    f"{index}-置信度：{round(line['score'], 2)}，文本：{line['text']}",
                    end="\\n\n" if line.get("end", "") == "\n" else "\n",
                )
                index += 1
        elif res["code"] == 100:
            print("图片中未识别出文字。")
        else:
            print(f"图片识别失败。错误码：{res['code']}，错误信息：{res['data']}")

    def __del__(self):
        self.exit()


class PPOCR_socket(PPOCR_pipe):
    """调用OCR（套接字模式）"""

    def __init__(self, exePath: str, modelsPath: str = None, argument: dict = None):
        """初始化识别器（套接字模式）。\n
        `exePath`: 识别器`PaddleOCR_json.exe`的路径。\n
        `modelsPath`: 识别库`models`文件夹的路径。若为None则默认识别库与识别器在同一目录下。\n
        `argument`: 启动参数，字典`{"键":值}`。参数说明见 https://github.com/hiroi-sora/PaddleOCR-json
        """
        # 处理参数
        if not argument:
            argument = {}
        if "port" not in argument:
            argument["port"] = 0  # 随机端口号
        if "addr" not in argument:
            argument["addr"] = "loopback"  # 本地环回地址

        # 处理输入的路径，可能为本地或远程路径
        self.__runningMode = self.__configureExePath(exePath)

        # 如果为本地路径：使用 PPOCR_pipe 来开启本地引擎进程
        if self.__runningMode == "local":
            super().__init__(self.exePath, modelsPath, argument)  # 父类构造函数
            self.__ENABLE_CLIPBOARD = super().isClipboardEnabled()
            # 再获取一行输出，检查是否成功启动服务器
            initStr = self.ret.stdout.readline().decode("utf-8", errors="ignore")
            if not self.ret.poll() == None:  # 子进程已退出，初始化失败
                raise Exception(f"Socket init fail.")
            if "Socket init completed. " in initStr:  # 初始化成功
                splits = initStr.split(":")
                self.ip = splits[0].split("Socket init completed. ")[1]
                self.port = int(splits[1])  # 提取端口号
                self.ret.stdout.close()  # 关闭管道重定向，防止缓冲区填满导致堵塞
                print(f"套接字服务器初始化成功。{self.ip}:{self.port}")
                return

        # 如果为远程路径：直接连接
        elif self.__runningMode == "remote":
            self.__ENABLE_CLIPBOARD = False
            # 发送一个空指令，检测远程服务器可用性
            testServer = self.runDict({})
            if testServer["code"] in [902, 903, 904]:
                raise Exception(f"Socket connection fail.")
            print(f"套接字服务器连接成功。{self.ip}:{self.port}")
            return

        # 异常
        self.exit()
        raise Exception(f"Socket init fail.")

    def isClipboardEnabled(self) -> bool:
        return self.__ENABLE_CLIPBOARD

    def getRunningMode(self) -> str:
        return self.__runningMode

    def runDict(self, writeDict: dict):
        """传入指令字典，发送给引擎进程。\n
        `writeDict`: 指令字典。\n
        `return`:  {"code": 识别码, "data": 内容列表或错误信息字符串}\n"""

        # 仅在本地模式下检查引擎进程
        if self.__runningMode == "local":
            # 检查子进程
            if not self.ret.poll() == None:
                return {"code": 901, "data": f"子进程已崩溃。"}

        # 通信
        writeStr = jsonDumps(writeDict, ensure_ascii=True, indent=None) + "\n"
        try:
            # 创建TCP连接
            clientSocket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            clientSocket.connect((self.ip, self.port))
            # 发送数据
            clientSocket.sendall(writeStr.encode())
            # 发送完所有数据，关闭我方套接字，之后只能从服务器读取数据
            clientSocket.shutdown(socket.SHUT_WR)
            # 接收数据
            resData = b""
            while True:
                chunk = clientSocket.recv(1024)
                if not chunk:
                    break
                resData += chunk
            getStr = resData.decode()
        except ConnectionRefusedError:
            return {"code": 902, "data": "连接被拒绝"}
        except TimeoutError:
            return {"code": 903, "data": "连接超时"}
        except Exception as e:
            return {"code": 904, "data": f"网络错误：{e}"}
        finally:
            clientSocket.close()  # 关闭连接
        # 反序列输出信息
        try:
            return jsonLoads(getStr)
        except Exception as e:
            return {
                "code": 905,
                "data": f"识别器输出值反序列化JSON失败。异常信息：[{e}]。原始内容：[{getStr}]",
            }

    def exit(self):
        """关闭引擎子进程"""
        # 仅在本地模式下关闭引擎进程
        if hasattr(self, "ret"):
            if self.__runningMode == "local":
                if not self.ret:
                    return
                try:
                    self.ret.kill()  # 关闭子进程
                except Exception as e:
                    print(f"[Error] ret.kill() {e}")
            self.ret = None

        self.ip = None
        self.port = None
        atexit.unregister(self.exit)  # 移除退出处理
        print("###  PPOCR引擎子进程关闭！")

    def __del__(self):
        self.exit()

    def __configureExePath(self, exePath: str) -> str:
        """处理识别器路径，自动区分本地路径和远程路径"""

        pattern = r"remote://(.*):(\d+)"
        match = re.search(pattern, exePath)
        try:
            if match:  # 远程模式
                self.ip = match.group(1)
                self.port = int(match.group(2))
                if self.ip == "any":
                    self.ip = "0.0.0.0"
                elif self.ip == "loopback":
                    self.ip = "127.0.0.1"
                return "remote"
            else:  # 本地模式
                self.exePath = exePath
                return "local"
        except:
            return None


def GetOcrApi(
    exePath: str, modelsPath: str = None, argument: dict = None, ipcMode: str = "pipe"
):
    """获取识别器API对象。\n
    `exePath`: 识别器`PaddleOCR_json.exe`的路径。\n
    `modelsPath`: 识别库`models`文件夹的路径。若为None则默认识别库与识别器在同一目录下。\n
    `argument`: 启动参数，字典`{"键":值}`。参数说明见 https://github.com/hiroi-sora/PaddleOCR-json\n
    `ipcMode`: 进程通信模式，可选值为套接字模式`socket` 或 管道模式`pipe`。用法上完全一致。
    """
    if ipcMode == "socket":
        return PPOCR_socket(exePath, modelsPath, argument)
    elif ipcMode == "pipe":
        return PPOCR_pipe(exePath, modelsPath, argument)
    else:
        raise Exception(
            f'ipcMode可选值为 套接字模式"socket" 或 管道模式"pipe" ，不允许{ipcMode}。'
        )