import os
import re
import json
import time
import shutil
import zipfile
from urllib.parse import urlparse
from uuid import uuid4
from PySide2.QtCore import QMutex
from typing import Dict

from umi_log import logger
from call_func import CallFunc
from .bottle import request, static_file, HTTPError
from .ocr_server import get_ocr_options
from ..ocr.output import Output
from ..mission.mission_doc import MissionDOC, DocSuf
from ..utils.utils import initConfigDict
from ..ocr.output.tools import getDataText

UPLOAD_DIR = "./temp_doc"  # 上传文件临时目录
TEMP_FILE_RETENTION_DURATION = 24  # 任务临时文件保留时长，小时
TEMP_FILE_CLEANUP_INTERVAL = 0.5  # 自动清理临时文件的间隔，小时


# 获取参数模板字典
def get_doc_options():
    opts = get_ocr_options(is_format=False)
    opts["tbpu.ignoreRangeStart"] = {
        "title": "忽略区域起始",
        "toolTip": "忽略区域生效的页数范围起始。从1开始。",
        "default": 1,
        "isInt": True,
    }
    opts["tbpu.ignoreRangeEnd"] = {
        "title": "忽略区域结束",
        "toolTip": "忽略区域生效的页数范围结束。可以用负数表示倒数第X页。",
        "default": -1,
        "isInt": True,
    }
    opts["pageRangeStart"] = {
        "title": "OCR页数起始",
        "toolTip": "OCR的页数范围起始。从1开始。",
        "default": 1,
        "isInt": True,
    }
    opts["pageRangeEnd"] = {
        "title": "OCR页数结束",
        "toolTip": "OCR的页数范围结束。可以用负数表示倒数第X页。",
        "default": -1,
        "isInt": True,
    }
    opts["pageList"] = {
        "title": "OCR页数列表",
        "toolTip": "数组，可指定单个或多个页数。例：[1,2,5]表示对第1、2、5页进行OCR。如果与页数范围同时填写，则 pageList 优先。",
        "default": [],
        "type": "var",
    }
    opts["password"] = {
        "title": "密码",
        "toolTip": "如果文档已加密，则填写文档密码。",
        "default": "",
    }
    opts["doc.extractionMode"] = {
        "title": "内容提取模式",
        "toolTip": "若一页文档既存在图片又存在文本，如何进行处理。",
        "optionsList": [
            ["mixed", "混合OCR/原文本"],
            ["fullPage", "整页强制OCR"],
            ["imageOnly", "仅OCR图片"],
            ["textOnly", "仅拷贝原有文本"],
        ],
    }
    opts = initConfigDict(opts)  # 格式化
    return opts


UPLOAD_DIR = os.path.abspath(UPLOAD_DIR)  # 路径转绝对
TEMP_FILE_RETENTION_DURATION *= 3600  # 小时转为秒
TEMP_FILE_CLEANUP_INTERVAL *= 3600


# 异常类
class DocUnitError(Exception):
    def __init__(self, data):
        self.data = data


# 单个任务单元
class _DocUnit:
    def __init__(
        self, dir_id, dir_path, origin_path, origin_name, origin_prefix, options
    ):
        # 提取文档信息
        doc_info = MissionDOC.getDocInfo(origin_path)
        if "error" in doc_info.keys():
            raise DocUnitError({"code": 201, "data": doc_info["error"]})

        # 补充缺失的默认参数
        default = get_doc_options()
        for key in default:
            if key not in options:
                options[key] = default[key]["default"]

        # 提取参数
        page_range = [options["pageRangeStart"], options["pageRangeEnd"]]  # 识别范围
        page_list = options["pageList"]  # 页数列表
        if page_list:  # 下标起始由1转为0
            page_list = [x - 1 for x in page_list]
        password = options["password"]  # 密码
        if not password and doc_info["is_encrypted"]:
            raise DocUnitError(
                {
                    "code": 202,
                    "data": "The doc is encrypted, please fill in the password.",
                }
            )

        # 从 options 中提取一些条目，组装 docArgd 作为 MissionDoc 任务参数字典
        prefixes = ["ocr.", "doc.", "tbpu."]  # 要提取的条目前缀
        doc_argd = {}
        for k, v in options.items():
            for prefix in prefixes:
                if k.startswith(prefix):
                    doc_argd[k] = v
                    break

        # 任务信息
        msnInfo = {
            "onStart": self._onStart,
            "onGet": self._onGet,
            "onEnd": self._onEnd,
            "argd": doc_argd,
        }

        # 提交任务
        self.msnID = ""
        msg = MissionDOC.addMission(
            msnInfo, origin_path, page_range, page_list, password
        )
        if not msg:
            raise DocUnitError({"code": 203, "data": "addMission unknow."})
        if msg.startswith("["):
            raise DocUnitError({"code": 204, "data": msg})
        page_list = msnInfo["pageList"]

        self.password = password
        self.dir_id = dir_id
        self.dir_path = dir_path
        self.origin_prefix = origin_prefix
        self.origin_name = origin_name
        self.origin_path = origin_path
        self.msnID = msg  # 任务ID
        self.results = {}  # 任务结果原始字典，键为页数
        self.pages_count = len(page_list)  # 任务总页数
        self.processed_count = 0  # 已处理的页数
        self.unread_list = []  # 未读的任务列表
        self.is_done = False  # 当前任务是否完成
        self.state = "waiting"  # 任务状态， waiting running success failure
        self.message = ""  # 如果任务失败，则记录失败信息
        self.start_timestamp = time.time()  # 开始时间戳
        self.end_timestamp = time.time()  # 任务结束的时间戳
        self._mutex = QMutex()  # 主锁

    # ========================= 【接口】 =========================

    # 获取结果
    def get_result(
        self,
        is_data=False,  # True 时返回识别内容data
        format="dict",  # 识别内容格式， "dict", "text"
        is_unread=False,  # True 时只返回未读过的识别内容
    ):
        self._mutex.lock()
        data = {
            "code": 100,
            "processed_count": self.processed_count,  # 已处理的数量
            "pages_count": self.pages_count,  # 总页数
            "is_done": self.is_done,  # 是否已结束
            "state": self.state,  # 任务状态
            "data": [],  # 结果
        }
        if self.state == "failure":
            data["message"] = self.message
        # 需要返回识别内容
        if is_data:
            datas = []
            # 增量式
            if is_unread:
                for page in self.unread_list:
                    datas.append(self.results[page])
                self.unread_list = []
            # 全量式
            else:
                for _, res in self.results.items():
                    datas.append(res)
            # 需要转为纯文本
            if format == "text":
                datas_text = ""
                for res in datas:
                    if res["code"] == 100:
                        datas_text += getDataText(res["data"])
                datas = datas_text
            data["data"] = datas
        self._mutex.unlock()
        return data

    # 获取文件
    def get_files(
        self,
        base_url,  # 下载基础url
        file_types=["pdfLayered"],  # 输出文件类型，可选：
        # txt, txtPlain, jsonl, csv, pdfLayered, pdfOneLayer
        ignore_blank=True,  # 忽略空白页数
    ):
        if not self.is_done:
            return {"code": 201, "data": f"{self.msnID} 任务尚未结束，无法获取文件"}
        if not self.state == "success":
            return {"code": 201, "data": f"{self.msnID} 任务处理失败，无法获取文件"}
        if not isinstance(file_types, list) or not isinstance(ignore_blank, bool):
            return {
                "code": 202,
                "data": f"参数类型错误： file_types={file_types} , ignore_blank={ignore_blank}",
            }

        # 删除旧的文件
        for filename in os.listdir(self.dir_path):
            file_path = os.path.join(self.dir_path, filename)
            if filename != self.origin_name and os.path.isfile(file_path):
                os.remove(file_path)

        # 准备参数
        startDatetime = time.strftime(  # 日期时间字符串（标准格式）
            r"%Y-%m-%d %H:%M:%S", time.localtime(self.start_timestamp)
        )
        outputArgd = {
            "outputDir": self.dir_path,  # 输出路径
            "outputDirType": "specify",
            "outputFileName": "[OCR]_" + self.origin_prefix,  # 输出文件名（前缀）
            "startDatetime": startDatetime,  # 开始日期
            "ignoreBlank": ignore_blank,  # 忽略空白页数
            "originPath": self.origin_path,  # 原始文件
            "password": self.password,  # 文档密码
        }

        # 创建输出器
        output = []
        try:
            for f in file_types:
                output.append(Output[f](outputArgd))
        except Exception as e:
            return {"code": 203, "data": f"初始化输出器失败。{e}"}

        # 输出
        for o in output:
            for _, res in self.results.items():
                try:
                    o.print(res)
                except Exception as e:
                    return {"code": 204, "data": f"输出失败：{o}\n{e}"}
            try:
                o.onEnd()  # 保存
            except Exception as e:
                return {"code": 205, "data": f"保存失败：{o}\n{e}"}

        # 收集新的文件
        download_paths = []
        for filename in os.listdir(self.dir_path):
            file_path = os.path.join(self.dir_path, filename)
            if filename != self.origin_name and os.path.isfile(file_path):
                download_paths.append(file_path)
        # 如果文件多，则打包zip
        if not download_paths:
            return {"code": 206, "data": "未找到生成的文件"}
        elif len(download_paths) == 1:
            download_name = os.path.basename(download_paths[0])
        else:
            download_name = f"[OCR]_{self.origin_prefix}.zip"
            zip_path = os.path.join(self.dir_path, download_name)
            # 将 download_list 中的所有文件打包为 zip
            try:
                with zipfile.ZipFile(zip_path, "w") as zipf:
                    for p in download_paths:
                        zipf.write(p, os.path.basename(p))
            except Exception as e:
                return {"code": 207, "data": f"无法打包zip：{e}"}

        # 组合下载地址
        url = f"{base_url}/api/doc/download/{self.dir_id}/{download_name}"

        return {"code": 100, "data": url, "name": download_name}

    # 清理任务
    def clear(self):
        # 停止任务
        if not self.is_done:
            MissionDOC.stopMissionList([self.msnID])
            time.sleep(0.1)  # 给一些时间收尾
        # 尝试删除目录。如果权限原因无法删除，说明OCR线程还在占用，等待OCR结束。
        for i in range(20):
            try:
                if os.path.exists(self.dir_path):
                    shutil.rmtree(self.dir_path)
            except PermissionError:
                time.sleep(1)

    # ========================= 【任务控制器的异步回调】 =========================

    def _onStart(self, msnInfo):  # 一个文档 开始
        self.state = "running"

    def _onGet(self, msnInfo, page, res):  # 一个文档的一页 获取结果
        page += 1
        res["page"] = page
        res["path"] = f"{self.origin_name} - {page}"
        res["fileName"] = f"{self.origin_name} - {page}"

        # 记录信息
        self._mutex.lock()
        self.results[page] = res
        self.processed_count += 1
        self.unread_list.append(page)
        self._mutex.unlock()

    def _onEnd(self, msnInfo, msg):  # 一个文档处理完毕
        # msg: [Success] [Warning] [Error]

        # 记录信息
        self._mutex.lock()
        self.is_done = True
        if msg == "[Success]":
            self.state = "success"
        else:
            self.state = "failure"
            self.message = msg
        self.end_timestamp = time.time()  # 刷新结束时间戳
        self._mutex.unlock()


# 管理所有任务单元
class _DocUnitManagerClass:
    def __init__(self):
        self.doc_units: Dict[str, _DocUnit] = {}

    # 添加一个任务单元
    def add(self, id: str, unit: _DocUnit):
        self.doc_units[id] = unit

    # 获取一个任务单元
    def get(self, id: str):
        if id not in self.doc_units:
            return None
        return self.doc_units[id]

    # 手动清理一个任务
    def clear(self, id: str):
        if id in self.doc_units:
            self.doc_units[id].clear()
            del self.doc_units[id]
            return True
        return False

    # 自动清理
    def auto_clear(self):
        # 清理超时的任务和文件
        if self.doc_units:
            now = time.time()  # 当前时间戳
            del_list = []  # 要清理的id
            for id, unit in self.doc_units.items():
                if now - unit.end_timestamp > TEMP_FILE_RETENTION_DURATION:
                    logger.info(f"超时自动清理 {id}")
                    unit.clear()  # 清理文件
                    del_list.append(id)
            for id in del_list:
                del self.doc_units[id]  # 清理任务对象
        # 计划下一次清理
        CallFunc.delay(self.auto_clear, TEMP_FILE_CLEANUP_INTERVAL)


_DocUnitManager = _DocUnitManagerClass()


# 路由函数
def init(UmiWeb):
    # 清空上传文件目录内容
    if os.path.exists(UPLOAD_DIR):
        shutil.rmtree(UPLOAD_DIR)
    os.makedirs(UPLOAD_DIR)
    # 启动自动清理循环
    _DocUnitManager.auto_clear()

    @UmiWeb.route("/api/doc/get_options")
    def _get_options_json():
        opts = get_doc_options()
        res = json.dumps(opts)
        return res

    """
    上传文档，方法：POST
    参数：文档内容
    返回值：
    成功： {"code": 100, "data": "任务id"}
    失败： {"code": 不是100的值, "data": "失败原因"}
    """

    @UmiWeb.route("/api/doc/upload", method="POST")
    def _upload():
        # 1. 获取上传文件
        upload = request.files.get("file")
        if not upload:
            return {"code": 101, "data": "[Error] No file was uploaded."}

        # 2. 获取文件名，检查文件后缀
        # origin_name = upload.filename
        # 将原始文件名转为合法文件名
        def filename_convert(raw_filename: str):
            # 去除前后的空格
            raw_filename = raw_filename.strip()
            # 定义非法字符
            illegal_chars = r'[\/:*?"<>|]'
            # 替换非法字符
            sanitized_filename = re.sub(illegal_chars, "_", raw_filename)
            # 限制文件名长度为255个字符
            max_length = 255
            if len(sanitized_filename) > max_length:
                sanitized_filename = sanitized_filename[:max_length]
            return sanitized_filename

        try:
            origin_name = filename_convert(upload.raw_filename)
            origin_prefix, ext = os.path.splitext(origin_name)
            ext = ext.lower()
        except Exception as e:
            return {"code": 102, "data": f"[Error] Unable to obtain raw_filename: {e}"}
        if ext not in DocSuf:
            return {
                "code": 103,
                "data": f"[Error] File extension '{ext}' is not allowed.",
            }

        # 3. 指定文件编号。创建对应目录，保存文件到 ./temp/dir_id/原文件名
        dir_id = str(uuid4())
        dir_path = os.path.join(UPLOAD_DIR, f"{dir_id}")
        dir_path = os.path.abspath(dir_path)  # 将路径转为绝对路径
        file_path = os.path.join(dir_path, origin_name)
        # 安全检测： file_path 是否在 UPLOAD_DIR 中
        if os.path.commonpath([UPLOAD_DIR]) != os.path.commonpath(
            [UPLOAD_DIR, file_path]
        ):
            return {"code": 104, "data": f"[Error] Unauthorized path"}

        try:
            if os.path.exists(dir_path):  # 如果目录存在，则删除该目录
                shutil.rmtree(dir_path)
            os.makedirs(dir_path)  # 重新创建目录
        except Exception as e:
            return {"code": 105, "data": f"[Error] Failed to create dir_id: {e}"}
        try:
            upload.save(file_path, overwrite=True)  # 保存文件
        except Exception as e:
            return {"code": 106, "data": f"[Error] Failed to save file: {e}"}

        # 4. 提取 options 参数
        options = request.forms.get("json")
        if options:
            try:
                options = json.loads(options)
            except Exception as e:
                shutil.rmtree(dir_path)
                return {
                    "code": 107,
                    "data": f"[Error] Invalid JSON format: {options} | {e}",
                }
        if not isinstance(options, dict):
            options = {}

        # 5. 构造任务对象
        try:
            doc_unit = _DocUnit(
                dir_id, dir_path, file_path, origin_name, origin_prefix, options
            )
            msnID = doc_unit.msnID
            _DocUnitManager.add(msnID, doc_unit)
            logger.info(f"添加 HTTP 文档任务: {origin_name}")
            return {"code": 100, "data": msnID}
        except DocUnitError as e:
            shutil.rmtree(dir_path)
            return e.data
        except Exception as e:
            shutil.rmtree(dir_path)
            return {"code": 108, "data": f"[Error] Failed to submit mission: {e}"}

    """
    获取结果，方法：POST
    json参数：
    "id"="",  # 任务ID
    "is_data"=False,  # True 时返回识别内容data
    "format"="dict",  # 识别内容格式， "dict", "text"
    "is_unread"=False,  # True 时只返回未读过的识别内容

    返回值： {}
    """

    @UmiWeb.route("/api/doc/result", method="POST")
    def _result():
        try:
            user_data = request.json
        except Exception as e:
            return {"code": 101, "data": f"请求无法解析为json。"}
        if not user_data or "id" not in user_data:
            return {"code": 102, "data": f"未填写id。"}
        msnID = user_data["id"]
        doc_unit = _DocUnitManager.get(msnID)
        if not doc_unit:
            return {"code": 103, "data": f"任务 {msnID} 不存在。"}
        is_data = user_data.get("is_data", False)
        format = user_data.get("format", "dict")
        is_unread = user_data.get("is_unread", True)
        return doc_unit.get_result(is_data, format, is_unread)

    """
    获取文件，方法：POST
    json参数：
    "id"="",  # 任务ID
    "file_types"=["pdfLayered"],  # 输出文件类型，可选：
    # ["txt", "txtPlain", "jsonl", "csv", "pdfLayered", "pdfOneLayer"]
    "ignore_blank"=True,  # 忽略空白页数

    返回值： {}
    """

    @UmiWeb.route("/api/doc/download", method="POST")
    def _download_build():
        try:
            user_data = request.json
        except Exception as e:
            return {"code": 101, "data": f"请求无法解析为json。"}
        if not user_data or "id" not in user_data:
            return {"code": 102, "data": f"未填写id。"}
        msnID = user_data["id"]
        doc_unit = _DocUnitManager.get(msnID)
        if not doc_unit:
            return {"code": 103, "data": f"任务 {msnID} 不存在。"}

        file_types = user_data.get("file_types", ["pdfLayered"])
        ignore_blank = user_data.get("ignore_blank", True)
        parsed_url = urlparse(request.url)
        base_url = f"{parsed_url.scheme}://{parsed_url.netloc}"
        return doc_unit.get_files(base_url, file_types, ignore_blank)

    # 下载文件
    @UmiWeb.route("/api/doc/download/<id>/<download_name>")
    def _download_get(id, download_name):
        dir = os.path.join(UPLOAD_DIR, id)
        path = os.path.join(dir, download_name)
        # 安全检测： path 是否在 UPLOAD_DIR 中
        if os.path.commonpath([UPLOAD_DIR]) != os.path.commonpath([UPLOAD_DIR, path]):
            raise HTTPError(103, "[Error] Unauthorized path.")
        return static_file(download_name, root=dir)

    # 清理任务
    @UmiWeb.route("/api/doc/clear/<id>")
    def _clear(id):
        flag = _DocUnitManager.clear(id)
        if flag:
            return {"code": 100, "data": "Success"}
        return {"code": 101, "data": f"{id} does not exist."}