# logic.py
import os
import sys

# 添加父目录到sys.path，便于导入onnxocr包
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import time
from pathlib import Path
from typing import Callable, List

import cv2
import numpy as np

from onnx_ocr.onnx_paddleocr import ONNXPaddleOcr, sav2Img

# 尝试导入pdf2image用于PDF转图片
try:
    from pdf2image import convert_from_path
except ImportError:
    convert_from_path = None

# 尝试导入pymupdf用于PDF转图片
try:
    import fitz  # pymupdf

    def pdf_to_images(pdf_path, dpi=200):
        """
        使用pymupdf将PDF每一页转为图片（numpy数组）
        """
        doc = fitz.open(pdf_path)
        images = []
        for page in doc:
            pix = page.get_pixmap(dpi=dpi)
            img = np.frombuffer(pix.samples, dtype=np.uint8)
            img = img.reshape((pix.height, pix.width, pix.n))
            if pix.n == 4:
                img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            images.append(img)
        return images
except ImportError:
    pdf_to_images = None


class OCRLogic:
    """
    OCR 业务逻辑主类，支持批量图片/PDF识别，多线程加速，模型热切换等
    """

    def __init__(self, status_callback: Callable[[str], None]):
        """
        初始化，传入状态回调函数用于UI进度提示
        """
        self.status_callback = status_callback
        # 默认初始化OCR模型
        self.model = ONNXPaddleOcr(use_angle_cls=True, use_gpu=False)

    def run(
        self,
        files: List[str],
        save_txt: bool,
        merge_txt: bool,
        output_img: bool = False,
        file_time_callback=None,
        pdf_progress_callback=None,
        max_workers: int = 4,
    ):
        """
        批量图片/PDF识别主入口，支持多线程加速
        files: 待识别文件路径列表
        save_txt: 是否保存txt
        merge_txt: 是否合并为一个txt
        output_img: 是否输出带框图片
        file_time_callback: 单文件识别耗时回调
        pdf_progress_callback: PDF页进度回调
        max_workers: 最大线程数，默认4
        """
        import concurrent.futures

        start_time = time.time()
        all_text = [None] * len(files)  # 用于顺序合并结果

        def process_one(idx_file):
            idx, file = idx_file
            ext = os.path.splitext(file)[1].lower()
            self.status_callback(
                f"正在处理: {os.path.basename(file)} ({idx + 1}/{len(files)})"
            )
            t0 = time.time()
            text = ""
            if ext == ".pdf":
                # PDF转图片后识别
                if pdf_to_images is None:
                    raise RuntimeError(
                        "未安装pymupdf库，无法处理PDF文件。请先安装pymupdf。"
                    )
                images = pdf_to_images(file, dpi=300)
                text = self._ocr_images(
                    images,
                    file,
                    save_txt,
                    merge_txt,
                    output_img=output_img,
                    is_pdf=True,
                    pdf_progress_callback=pdf_progress_callback,
                    max_workers=max_workers,
                )
            else:
                # 普通图片识别，兼容中文路径
                try:
                    if file.lower().endswith(".bmp"):
                        img = cv2.imdecode(
                            np.fromfile(file, dtype=np.uint8), cv2.IMREAD_COLOR
                        )
                    else:
                        with open(file, "rb") as fimg:
                            img_array = np.frombuffer(fimg.read(), np.uint8)
                        img = cv2.imdecode(img_array, cv2.IMREAD_COLOR)
                except Exception as e:
                    self.status_callback(f"图片读取失败: {file}，错误: {e}")
                    if file_time_callback:
                        file_time_callback(idx, 0)
                    return (idx, "")
                if img is None:
                    self.status_callback(f"文件无法读取或不是有效图片: {file}")
                    if file_time_callback:
                        file_time_callback(idx, 0)
                    return (idx, "")
                text = self._ocr_image(img, file, save_txt, output_img=output_img)
            t1 = time.time()
            if file_time_callback:
                file_time_callback(idx, t1 - t0)
            self.status_callback(f"{os.path.basename(file)} 识别用时: {t1 - t0:.2f} 秒")
            if len(files) > 1:
                avg = (t1 - start_time) / (idx + 1)
                self.status_callback(
                    f"已完成 {idx + 1}/{len(files)}，平均单张用时: {avg:.2f} 秒"
                )
            return (idx, text)

        # 多线程处理所有文件，结果按索引回填，保证顺序
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [
                executor.submit(process_one, (idx, file))
                for idx, file in enumerate(files)
            ]
            for future in concurrent.futures.as_completed(futures):
                idx, text = future.result()
                all_text[idx] = text
        # 合并写入txt
        if save_txt and merge_txt and len(files) > 1:
            out_dir = self._get_output_dir(files[0])
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            out_txt = os.path.join(out_dir, f"merged_ocr_{timestamp}.txt")
            with open(out_txt, "w", encoding="utf-8") as f:
                for text in all_text:
                    if text:
                        f.write(text)
                        f.write("\n\n")
        elapsed = time.time() - start_time
        if files:
            out_dir = self._get_output_dir(files[0])
            self.status_callback(
                f"识别完成，总耗时：{elapsed:.2f}秒，文件保存在：{out_dir}"
            )
        else:
            self.status_callback(f"识别完成，总耗时：{elapsed:.2f}秒")

    def _ocr_images(
        self,
        images,
        pdf_path,
        save_txt,
        merge_txt,
        output_img=False,
        is_pdf=False,
        pdf_progress_callback=None,
        max_workers: int = 4,
    ):
        """
        PDF转图片后，批量图片识别，支持多线程加速
        images: PDF每页图片（numpy数组）
        pdf_path: 原PDF路径
        save_txt: 是否保存txt
        merge_txt: 是否合并txt（未用）
        output_img: 是否输出带框图片
        pdf_progress_callback: 页进度回调
        max_workers: 最大线程数，默认4
        """
        import concurrent.futures

        out_dir = self._get_output_dir(pdf_path)
        pdf_text = [None] * len(images)
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        total = len(images)

        def process_page(i_img):
            i, img = i_img
            img_cv = cv2.cvtColor(np.array(img), cv2.COLOR_RGB2BGR)
            result = self.model.ocr(img_cv)
            if output_img:
                out_img_path = os.path.join(
                    out_dir, f"{Path(pdf_path).stem}_page{i + 1}_ocr.jpg"
                )
                sav2Img(img_cv, result, name=out_img_path)
            page_text = self._result_to_text(result)
            return (i, page_text)

        # 多线程识别每一页，结果按页码顺序合并
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
            futures = [
                executor.submit(process_page, (i, img)) for i, img in enumerate(images)
            ]
            for future in concurrent.futures.as_completed(futures):
                i, page_text = future.result()
                pdf_text[i] = page_text
                if pdf_progress_callback:
                    pdf_progress_callback(i + 1, total)
        if save_txt:
            txt_path = os.path.join(
                out_dir, f"{Path(pdf_path).stem}_ocr_{timestamp}.txt"
            )
            with open(txt_path, "w", encoding="utf-8") as f:
                f.write("\n\n".join(pdf_text))
        return "\n\n".join(pdf_text)

    def _ocr_image(self, img, img_path, save_txt, output_img=False):
        """
        单张图片OCR识别，支持保存txt和输出带框图片
        """
        out_dir = self._get_output_dir(img_path)
        result = self.model.ocr(img)
        if output_img:
            out_img_path = os.path.join(out_dir, f"{Path(img_path).stem}_ocr.jpg")
            sav2Img(img, result, name=out_img_path)
        text = self._result_to_text(result)
        if save_txt:
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            txt_path = os.path.join(
                out_dir, f"{Path(img_path).stem}_ocr_{timestamp}.txt"
            )
            with open(txt_path, "w", encoding="utf-8") as f:
                f.write(text)
        return text

    def _result_to_text(self, result):
        """
        将OCR识别结果结构化为纯文本，兼容只检测无识别内容的情况
        """
        # 健壮性检查，防止result为空或结构异常
        if (
            not result
            or not isinstance(result, list)
            or not result[0]
            or not isinstance(result[0], list)
        ):
            return "[未检测到内容]"
        lines = []
        for box in result[0]:
            # 兼容只检测无识别内容的情况
            if (
                isinstance(box, list)
                and len(box) == 2
                and isinstance(box[1], (list, tuple))
                and len(box[1]) >= 1
            ):
                lines.append(str(box[1][0]))
            elif isinstance(box, list) and (
                isinstance(box[0], (list, tuple)) or isinstance(box[0], float)
            ):
                # 只有检测框，无识别内容
                lines.append("[未识别] " + str(box))
            else:
                lines.append(str(box))
        return "\n".join(lines)

    def _get_output_dir(self, file_path):
        """
        获取输出目录，自动创建
        """
        base_dir = os.path.dirname(file_path)
        out_dir = os.path.join(base_dir, "Output_OCR")
        os.makedirs(out_dir, exist_ok=True)
        return out_dir

    def set_model(self, model_name, use_gpu=False):
        """
        切换OCR模型，支持多模型热切换，所有模型统一用ppocrv5字典
        use_gpu: 是否启用GPU
        """
        import os

        base_model_dir = os.path.abspath(
            os.path.join(os.path.dirname(__file__), "..", "onnxocr", "models")
        )
        model_map = {
            "PP-OCRv5": "ppocrv5",
            "PP-OCRv4": "ppocrv4",
            "ch_ppocr_server_v2.0": "ch_ppocr_server_v2.0",
        }
        model_dir = model_map.get(model_name, "ppocrv5")
        model_path = os.path.join(base_model_dir, model_dir)
        det_model_dir = os.path.join(model_path, "det", "det.onnx")
        cls_model_dir = os.path.join(model_path, "cls", "cls.onnx")
        rec_char_dict_path = os.path.join(base_model_dir, "ppocrv5", "ppocrv5_dict.txt")
        rec_model_dir = (
            os.path.join(model_path, "rec", "rec.onnx")
            if os.path.exists(os.path.join(model_path, "rec", "rec.onnx"))
            else None
        )
        ocr_kwargs = dict(
            use_angle_cls=True,
            use_gpu=use_gpu,  # 关键：传递GPU参数
            det_model_dir=det_model_dir,
            cls_model_dir=cls_model_dir,
            rec_char_dict_path=rec_char_dict_path,
        )
        if rec_model_dir and os.path.exists(rec_model_dir):
            ocr_kwargs["rec_model_dir"] = rec_model_dir
        try:
            self.model = ONNXPaddleOcr(**ocr_kwargs)
            if use_gpu:
                try:
                    providers = (
                        self.model.session.get_providers()
                        if hasattr(self.model, "session")
                        else []
                    )
                    if not any("CUDA" in p for p in providers):
                        msg = "未检测到可用GPU，已自动切换为CPU推理。请检查CUDA/cuDNN环境配置。"
                        if hasattr(self, "ui_ref") and hasattr(
                            self.ui_ref, "update_gpu_status"
                        ):
                            self.ui_ref.update_gpu_status(msg)
                        if hasattr(self, "status_callback"):
                            self.status_callback(
                                "[警告] 未检测到可用GPU，已切换为CPU推理。请检查CUDA/cuDNN环境配置。"
                            )
                except Exception:
                    msg = "检测GPU状态时发生异常，可能未正确安装CUDA/cuDNN或onnxruntime-gpu。已自动切换为CPU推理。"
                    if hasattr(self, "ui_ref") and hasattr(
                        self.ui_ref, "update_gpu_status"
                    ):
                        self.ui_ref.update_gpu_status(msg)
                    if hasattr(self, "status_callback"):
                        self.status_callback(
                            "[警告] GPU检测异常，已切换为CPU推理。请检查CUDA/cuDNN环境配置。"
                        )
        except Exception as e:
            if use_gpu:
                msg = f"GPU初始化失败，已自动切换为CPU。请检查CUDA/cuDNN环境配置。错误信息: {e}"
                if hasattr(self, "ui_ref") and hasattr(
                    self.ui_ref, "update_gpu_status"
                ):
                    self.ui_ref.update_gpu_status(msg)
                if hasattr(self, "status_callback"):
                    self.status_callback(
                        "[警告] GPU初始化失败，已切换为CPU推理。请检查CUDA/cuDNN环境配置。"
                    )
                ocr_kwargs["use_gpu"] = False
                self.model = ONNXPaddleOcr(**ocr_kwargs)
            else:
                raise
