# -*- coding: utf-8 -*-
"""
OCR 引擎模块
封装 PaddleOCR，提供统一的 OCR 接口
"""

import os
from pathlib import Path

# 在导入 PaddleOCR 之前设置环境变量
# 解决 Windows 中文用户名路径问题
_PROJECT_ROOT = Path(__file__).parent.parent
_MODELS_DIR = _PROJECT_ROOT / "models"
_MODELS_DIR.mkdir(exist_ok=True)
os.environ["PADDLEOCR_HOME"] = str(_MODELS_DIR)

import numpy as np
from typing import List, Optional, Any
from dataclasses import dataclass
from paddleocr import PaddleOCR

from utils.config import OCRConfig


@dataclass
class TextBlock:
    """
    文本块数据结构
    表示 OCR 识别出的单个文本区域

    Attributes:
        text: 识别出的文本内容
        confidence: 置信度 (0.0 ~ 1.0)
        bbox: 边界框，4 个点的坐标 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
        bbox_offset: ROI 偏移量，用于还原到原图坐标
    """
    text: str
    confidence: float
    bbox: List[List[float]]
    bbox_offset: tuple = (0, 0)

    @property
    def bbox_with_offset(self) -> List[List[float]]:
        """获取带偏移的边界框（还原到原图坐标）"""
        offset_x, offset_y = self.bbox_offset
        return [[p[0] + offset_x, p[1] + offset_y] for p in self.bbox]

    @property
    def center(self) -> tuple:
        """获取文本块中心点"""
        x_coords = [p[0] for p in self.bbox]
        y_coords = [p[1] for p in self.bbox]
        return (sum(x_coords) / 4, sum(y_coords) / 4)

    @property
    def width(self) -> float:
        """获取文本块宽度"""
        x_coords = [p[0] for p in self.bbox]
        return max(x_coords) - min(x_coords)

    @property
    def height(self) -> float:
        """获取文本块高度"""
        y_coords = [p[1] for p in self.bbox]
        return max(y_coords) - min(y_coords)

    def to_dict(self) -> dict:
        """转换为字典格式"""
        return {
            "text": self.text,
            "confidence": self.confidence,
            "bbox": self.bbox,
            "bbox_with_offset": self.bbox_with_offset,
            "center": self.center,
            "width": self.width,
            "height": self.height
        }


class OCREngine:
    """
    OCR 引擎类
    封装 PaddleOCR，提供简洁的 OCR 调用接口
    """

    def __init__(self, config: OCRConfig):
        """
        初始化 OCR 引擎

        Args:
            config: OCR 配置
        """
        self._config = config
        self._ocr: Optional[PaddleOCR] = None

    def initialize(self) -> None:
        """
        初始化 PaddleOCR 实例
        延迟初始化，避免在导入时加载模型
        适配 PaddleOCR 2.x API
        """
        if self._ocr is not None:
            return

        # 构建参数
        params = {
            "lang": self._config.lang,
            "use_angle_cls": self._config.use_angle_cls,
            "use_gpu": self._config.use_gpu,
            "det_db_thresh": self._config.det_db_thresh,
            "det_db_box_thresh": self._config.det_db_box_thresh,
            "drop_score": self._config.drop_score,
            "show_log": self._config.show_log
        }

        # 如果指定了模型目录，则使用自定义路径（解决中文路径问题）
        if self._config.det_model_dir:
            params["det_model_dir"] = self._config.det_model_dir
        if self._config.rec_model_dir:
            params["rec_model_dir"] = self._config.rec_model_dir
        if self._config.cls_model_dir:
            params["cls_model_dir"] = self._config.cls_model_dir

        # PaddleOCR 2.x API
        self._ocr = PaddleOCR(**params)

    def recognize(
        self,
        image: np.ndarray,
        roi_offset: tuple = (0, 0)
    ) -> List[TextBlock]:
        """
        对图像进行 OCR 识别

        Args:
            image: 输入图像 (numpy array, BGR 或灰度图)
            roi_offset: ROI 偏移量 (x, y)，用于还原坐标

        Returns:
            识别结果列表
        """
        # 确保引擎已初始化
        if self._ocr is None:
            self.initialize()

        # 执行 OCR (PaddleOCR 2.x API)
        result = self._ocr.ocr(image, cls=self._config.use_angle_cls)

        # 解析结果
        text_blocks: List[TextBlock] = []

        # PaddleOCR 返回格式: [[line1, line2, ...]] 或 None
        if result is None or len(result) == 0:
            return text_blocks

        # 遍历每一行结果
        for line in result:
            if line is None:
                continue
            for item in line:
                if item is None or len(item) < 2:
                    continue

                bbox = item[0]  # [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
                text_info = item[1]  # (text, confidence)

                if len(text_info) < 2:
                    continue

                text = text_info[0]
                confidence = float(text_info[1])

                # 过滤低置信度结果
                if confidence < self._config.drop_score:
                    continue

                text_block = TextBlock(
                    text=text,
                    confidence=confidence,
                    bbox=bbox,
                    bbox_offset=roi_offset
                )
                text_blocks.append(text_block)

        return text_blocks

    def recognize_batch(
        self,
        images: List[np.ndarray]
    ) -> List[List[TextBlock]]:
        """
        批量 OCR 识别

        Args:
            images: 输入图像列表

        Returns:
            每张图像的识别结果列表
        """
        return [self.recognize(img) for img in images]

    @property
    def config(self) -> OCRConfig:
        """获取当前配置"""
        return self._config

    def update_config(self, **kwargs) -> None:
        """
        更新配置并重新初始化引擎

        Args:
            **kwargs: 要更新的配置项
        """
        for key, value in kwargs.items():
            if hasattr(self._config, key):
                setattr(self._config, key, value)

        # 重新初始化
        self._ocr = None
        self.initialize()