import cv2
import numpy as np
from XZGUtil.ImagePositioning import window_rectangle
from rapidocr import RapidOCR

# 全局初始化引擎（避免重复创建）
engine = RapidOCR()


def preprocess_image(image):
    """
    图片预处理：增强对比度+膨胀，减少OCR文本粘连
    :param image: 输入截图（numpy数组，RGB格式）
    :return: 预处理后的RGB图片
    """
    if image is None:
        return image
    # 转为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    # 自适应二值化增强对比度
    thresh = cv2.adaptiveThreshold(
        gray, 255,
        cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY,
        blockSize=11,
        C=2
    )
    # 轻微膨胀，分隔粘连文本
    kernel = np.ones((1, 1), np.uint8)
    dilated = cv2.dilate(thresh, kernel, iterations=1)
    # 转回RGB格式（适配RapidOCR）
    return cv2.cvtColor(dilated, cv2.COLOR_GRAY2RGB)


def split_sticky_text(original_item, split_chars=["  ", " ", "\t"]):
    """
    拆分粘连文本，均分坐标（兜底解决文本粘连问题）
    :param original_item: 原始OCR结果字典
    :param split_chars: 粘连分隔符（双空格/单空格/制表符）
    :return: 拆分后的结果列表
    """
    original_text = original_item['text']
    if not original_text:
        return [original_item]

    # 遍历分隔符，找到第一个匹配的拆分符
    split_char = None
    for char in split_chars:
        if char in original_text:
            split_char = char
            break
    if split_char is None:
        return [original_item]

    # 拆分文本并过滤空值
    text_parts = [part.strip() for part in original_text.split(split_char) if part.strip()]
    if len(text_parts) <= 1:
        return [original_item]

    # 计算每个子文本的坐标（按长度均分原始矩形）
    original_box = original_item['box']
    total_width = original_box[2][0] - original_box[0][0]  # 原始矩形宽度
    avg_width = total_width / len(text_parts)
    split_results = []

    for i, text_part in enumerate(text_parts):
        # 计算子文本的矩形框（高度不变，宽度均分）
        start_x = original_box[0][0] + i * avg_width
        end_x = start_x + avg_width
        sub_box = [
            [int(start_x), original_box[0][1]],
            [int(end_x), original_box[1][1]],
            [int(end_x), original_box[2][1]],
            [int(start_x), original_box[3][1]]
        ]
        # 重新计算子文本中心点
        x1, y1 = sub_box[0]
        x3, y3 = sub_box[2]
        w = x3 - x1
        h = y3 - y1
        sub_center = [x1 + w // 2, y1 + h // 2]

        # 构建拆分后的结果
        split_results.append({
            'text': text_part,
            'scores': original_item['scores'],
            'box': sub_box,
            'center': sub_center
        })
    return split_results


def ocr_endpoint(win, word=None, fuzzy=False):
    """
    优化版OCR识别：解决文本粘连 + 完善匹配逻辑
    :param win: 窗口对象（含rectangle()方法，返回left/top属性）
    :param word: 匹配目标词（None时返回所有结果）
    :param fuzzy: 是否模糊匹配（仅word不为None时生效）
    :return: 匹配结果列表（屏幕绝对坐标，整数格式）
    """
    # 1. 截取窗口截图
    screenshot = window_rectangle(win)

    # 2. 解析窗口屏幕坐标
    window_rect = win.rectangle()
    win_left = window_rect.left
    win_top = window_rect.top

    # 3. 输入校验
    if not isinstance(screenshot, np.ndarray):
        raise ValueError("screenshot必须是numpy数组！")
    if screenshot.ndim != 3 or screenshot.shape[2] not in (3, 4):
        raise ValueError("screenshot必须是HWC格式的RGB/RGBA数组！")

    # 4. 截图预处理（解决文本粘连）+ 转RGB
    img_np = screenshot[:, :, :3]
    processed_img = preprocess_image(img_np)  # 预处理减少粘连

    # 5. 执行OCR识别
    ocr_result_obj = engine(processed_img)
    if not hasattr(ocr_result_obj, 'boxes') or len(ocr_result_obj.boxes) == 0:
        return []

    # 6. 遍历OCR结果，处理坐标+拆分粘连文本
    all_processed_items = []
    for box, text, score in zip(ocr_result_obj.boxes, ocr_result_obj.txts, ocr_result_obj.scores):
        # 基础规范化
        text_str = str(text).strip()
        box_list = box.tolist() if hasattr(box, 'tolist') else box
        score_float = float(score)

        # 转换为屏幕绝对坐标
        screen_box = []
        for (x, y) in box_list:
            screen_x = int(x) + win_left
            screen_y = int(y) + win_top
            screen_box.append([screen_x, screen_y])

        # 计算中心点
        x1, y1 = screen_box[0]
        x3, y3 = screen_box[2]
        w = x3 - x1
        h = y3 - y1
        screen_center = [x1 + w // 2, y1 + h // 2]

        # 构建原始结果项
        original_item = {
            'text': text_str,
            'scores': score_float,
            'box': screen_box,
            'center': screen_center
        }

        # 拆分粘连文本（核心解决“全部  文章”这类问题）
        split_items = split_sticky_text(original_item)
        all_processed_items.extend(split_items)

    # 7. 匹配逻辑优化（区分word=None/精确/模糊）
    final_results = []
    for item in all_processed_items:
        if word is None:
            # word为None时返回所有结果
            final_results.append(item)
        else:
            if item['text'] == word:
                # 精确匹配优先
                final_results.append(item)
            elif fuzzy and word in item['text']:
                # 模糊匹配（仅开启fuzzy时生效）
                final_results.append(item)

    return final_results
