# 双层可搜索 searchable pdf
# https://github.com/pymupdf/PyMuPDF/discussions/2299

import os
import fitz  # PyMuPDF

from umi_log import logger
from .output import Output


class OutputPdfLayered(Output):
    def __init__(self, argd):
        self.dir = argd["outputDir"]  # 输出路径（文件夹）
        self.originPath = argd["originPath"]  # 原始文件路径
        self.fileName = argd["outputFileName"]  # 文件名
        self.password = argd["password"]  # 密码
        self.outputPath = f"{self.dir}/{self.fileName}.layered.pdf"  # 输出路径
        self.pdf = None
        self.existentPages = []  # 已处理的页数
        self.isInsertFont = False  # 是否有字体嵌入
        self.opacity = 0  # 文本透明度为0
        try:
            self.font = fitz.Font("cjk")  # 字体
        except Exception as e:
            raise Exception(f"Failed to load cjk font. {e}\n无法加载cjk字体。")
        try:
            self.pdf = self._getPDF(self.originPath)  # 加载pymupdf对象
        except Exception as e:
            raise Exception(
                f"Failed to load doc file. {e}\n无法加载文档。\n{self.originPath}"
            )

    # 获取pdf文档对象，或将其它类型的文档转为PDF对象
    def _getPDF(self, path):
        # https://github.com/pymupdf/PyMuPDF-Utilities/blob/master/examples/convert-document/convert.py
        doc = fitz.open(path)
        # 如果已加密，则尝试解密
        if doc.is_encrypted and not doc.authenticate(self.password):
            raise Exception(
                f'The document is encrypted, and the password "{self.password}" is incorrect.\n文档已加密，输入密码不正确。'
            )
        if doc.is_pdf:
            return doc
        b = doc.convert_to_pdf()  # 转换为PDF格式的二进制数据
        pdf = fitz.open("pdf", b)  # 创建PDF文档对象
        try:
            pdf.set_toc(doc.get_toc())  # 复制原始文档的目录
        except Exception:
            logger.warning("pdf.set_toc error", exc_info=True, stack_info=True)
        # 复制原始文档的元数据（如作者、标题等）
        meta = doc.metadata
        if not meta["producer"]:
            meta["producer"] = "Umi-OCR & PyMuPDF v" + fitz.VersionBind
        if not meta["creator"]:
            meta["creator"] = "Umi-OCR & PyMuPDF PDF converter"
        pdf.set_metadata(meta)
        # 复制原始文档的链接
        for pinput in doc:
            links = pinput.get_links()
            pout = pdf[pinput.number]
            for link in links:
                if link["kind"] == fitz.LINK_NAMED:  # 不处理 named links
                    continue
                pout.insert_link(link)  # 写入新文档
        doc.close()  # 释放原文档
        return pdf

    # 计算填满宽和高的一行字体大小
    def _calculateFontSize(self, text, w, h):
        if h > w:  # 竖排转为横排计算
            w, h = h, w
        fontsize = round(h)  # 字体大小初值，假设为行高
        minSize = 5  # 大小下限
        getLen = lambda text, s: self.font.text_length(text, fontsize=s)
        while getLen(text, fontsize) > w and fontsize >= minSize:
            fontsize -= 1  # 尝试减小字体，直到行宽刚好小于界限
        while getLen(text, fontsize) < w:
            fontsize += 1  # 尝试增大字体，直到行宽刚好超过界限
        while getLen(text, fontsize) > w and fontsize >= minSize:
            fontsize -= 0.1  # 再次减小字体，将精度提升到 0.1
        return fontsize

    def print(self, res):  # 输出图片结果
        if not self.pdf:
            logger.error("self.pdf 未初始化。")
            return
        pno = res["page"] - 1  # 当前页数
        self.existentPages.append(pno)  # 记录已处理的页面
        if not res["code"] == 100:
            return  # 忽略空白

        page = self.pdf[pno]  # 当前页对象
        page.clean_contents()  # 内容流清理、语法更正，减少错误
        protation = page.rotation  # 获取页面旋转角度
        isInsertFont = False  # 当前是否进行过字体注入
        # 插入文本，用shape.insert_text（可编辑）或page.insert_text（不可编辑）
        for tb in res["data"]:
            if self.opacity == 0 and "from" in tb and tb["from"] == "text":
                continue  # 双层（透明文字）模式下，跳过直接提取的文本，只写入OCR文本
            if not isInsertFont:  # 页面插入字体
                self.isInsertFont = isInsertFont = True
                page.insert_font(fontname="cjk", fontbuffer=self.font.buffer)
            text = tb["text"]
            box = tb["box"]
            x0, y0 = box[0]
            x2, y2 = box[2]
            w = x2 - x0
            h = y2 - y0
            fontsize = self._calculateFontSize(text, w, h)
            # 插入点的 旋转后的坐标
            point = fitz.Point(x0, y2) * page.derotation_matrix
            page.insert_text(
                point,
                text,
                fontsize,
                fontname="cjk",
                rotate=protation,  # 文本角度设定
                stroke_opacity=self.opacity,  # 描边透明度
                fill_opacity=self.opacity,  # 填充（字体）透明度
            )

    def onEnd(self):  # 结束时保存。
        if not self.pdf:
            return
        # 删除未处理的页数
        for i in range(len(self.pdf) - 1, -1, -1):
            if i not in self.existentPages:
                self.pdf.delete_page(i)
        logger.info(f"保存{len(self.pdf)}页PDF：{self.outputPath}")
        if self.isInsertFont:  # 有任意页面嵌入字体，则构建字体子集
            try:  # 对于部分PDF，如用txt直接打印的，构建字体子集会失败。
                self.pdf.subset_fonts()  # 构建字体子集，减小文件大小。需要 fontTools 库
            except Exception:  # TODO: 失败原因？可能文件中实际并没有字体？
                logger.error("构建字体子集失败。", exc_info=True, stack_info=True)
            # 保存：压缩并进行3级垃圾回收。等同 ez_save
            self.save(self.pdf, self.outputPath, deflate=True, garbage=3)
        else:
            # 无嵌入字体，则直接保存，不压缩
            self.save(self.pdf, self.outputPath)

    def save(self, pdf, path, **options):  # 保存并关闭 pdf 对象
        try:
            # 尝试保存到指定路径
            pdf.save(path, **options)
        except Exception:
            # 保存失败，尝试保存到 ".temp" 路径
            tempPath = self.outputPath + ".temp"
            logger.warning(f"保存PDF失败。 path: {path}", exc_info=True)
            try:
                pdf.save(tempPath, **options)
                pdf.close()
            except Exception as e1:
                logger.error(
                    f"保存PDF到临时路径失败。 tempPath: {tempPath}", exc_info=True
                )
                raise Exception(f"[Error] Unable to save PDF to [{tempPath}]: {e1}")
            # 已保存到 .temp 并 close 原对象，尝试替换文件
            try:
                if os.path.exists(path):
                    os.remove(path)
                os.rename(tempPath, path)
            except Exception as e2:
                logger.warning(
                    f"保存PDF文件替换失败。保存到临时文件: {tempPath}", exc_info=True
                )

                raise Exception(
                    f"[Warning] Unable to save PDF: [{path}]. Exception: {e2}. Saved to temporary path: [{tempPath}]."
                )
        else:  # 正常结束
            pdf.close()