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

# 无参数初始化引擎（适配所有RapidOCR版本，避免参数报错）
engine = RapidOCR()

def ocr_endpoint(win, word=None, fuzzy=False):
    """
    保留原方法名，适配逻辑：
    - 不传word：返回全部OCR结果（屏幕绝对坐标）
    - 传word：精准返回关键词坐标（解决文本粘连）
    :param win: 窗口对象（含rectangle()方法，返回left/top属性）
    :param word: 匹配目标词（None时返回所有结果）
    :param fuzzy: 是否模糊匹配（仅word不为None时生效）
    :return: OCR结果列表（屏幕绝对坐标，整数格式）
    """
    # 1. 截取窗口截图
    screenshot = window_rectangle(win)
    # 获取窗口屏幕坐标
    window_rect = win.rectangle()
    win_left = window_rect.left
    win_top = window_rect.top

    # 输入校验
    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数组！")

    # 2. 预处理：增强对比度，减少文本粘连（替代参数优化）
    img_np = screenshot[:, :, :3]
    gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY)
    kernel = np.ones((2, 2), np.uint8)
    opening = cv2.morphologyEx(gray, cv2.MORPH_OPEN, kernel)
    processed_gray = cv2.adaptiveThreshold(
        opening, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
        cv2.THRESH_BINARY, blockSize=15, C=3
    )
    processed_img = cv2.cvtColor(processed_gray, cv2.COLOR_GRAY2RGB)

    # 3. 执行OCR（兼容新旧版返回格式）
    ocr_result_obj = engine(processed_img)
    # 解析boxes/txts/scores（兼容新旧版）
    if isinstance(ocr_result_obj, tuple):
        # 旧版：返回(结果列表, 耗时)
        ocr_result = ocr_result_obj[0]
        if not ocr_result:
            return []
        boxes = [item[0] for item in ocr_result]
        txts = [item[1] for item in ocr_result]
        scores = [item[2] for item in ocr_result]
    else:
        # 新版：RapidOCROutput对象
        if not hasattr(ocr_result_obj, 'boxes') or len(ocr_result_obj.boxes) == 0:
            return []
        boxes = ocr_result_obj.boxes
        txts = ocr_result_obj.txts
        scores = ocr_result_obj.scores

    # 4. 处理所有OCR结果，转换为屏幕绝对坐标
    all_processed_items = []
    for box, text, score in zip(boxes, txts, scores):
        text_str = str(text).strip()
        score_float = float(score)
        # 转换为屏幕绝对坐标
        box_list = box.tolist() if hasattr(box, 'tolist') else box
        screen_box = []
        for (x, y) in box_list:
            screen_box.append([int(x) + win_left, int(y) + win_top])
        # 计算中心点
        x1, y1 = screen_box[0]
        x3, y3 = screen_box[2]
        w = x3 - x1
        h = y3 - y1
        screen_center = [x1 + w // 2, y1 + h // 2]
        # 构建基础结果项
        base_item = {
            'text': text_str,
            'scores': score_float,
            'box': screen_box,
            'center': screen_center
        }
        all_processed_items.append(base_item)

    # 5. 按关键词逻辑返回结果
    if word is None:
        # 不传word：返回全部OCR结果
        return all_processed_items
    else:
        # 传word：精准匹配关键词（解决粘连）
        matched_results = []
        for item in all_processed_items:
            text_str = item['text']
            if not text_str:
                continue
            # 精确匹配
            if text_str == word:
                matched_results.append(item)
            # 模糊匹配（仅开启fuzzy时）
            elif fuzzy and word in text_str:
                # 字符级切割，计算关键词精准坐标
                total_width = item['box'][2][0] - item['box'][0][0]
                char_width = total_width / len(text_str)
                start_idx = text_str.index(word)
                end_idx = start_idx + len(word)
                # 重新计算关键词精准坐标
                kw_x1 = item['box'][0][0] + int(start_idx * char_width)
                kw_x3 = item['box'][0][0] + int(end_idx * char_width)
                kw_y1 = item['box'][0][1]
                kw_y3 = item['box'][2][1]
                # 替换为关键词精准坐标
                matched_item = item.copy()
                matched_item['text'] = word
                matched_item['keyword_box'] = [
                    [kw_x1, kw_y1],
                    [kw_x3, kw_y1],
                    [kw_x3, kw_y3],
                    [kw_x1, kw_y3]
                ]
                matched_item['keyword_center'] = [
                    kw_x1 + (kw_x3 - kw_x1) // 2,
                    kw_y1 + (kw_y3 - kw_y1) // 2
                ]
                matched_results.append(matched_item)
        return matched_results
