import sys
from typing import Any, Dict, List, Tuple, cast

import cv2
import numpy as np

# 确保 mineru 库可导入
mineru_path = str((__file__ and __file__) and __file__)
# 使用已有 mineru_adapter 中的路径追加逻辑
from pathlib import Path
mineru_root = Path(__file__).parents[4] / "mineru"
if str(mineru_root) not in sys.path:
    sys.path.insert(0, str(mineru_root))

from mineru.model.table.rec.unet_table.main import UnetTableModel


class MinerUWiredTableRecognizer:
    """有线表格识别封装：裁剪+放大→UNet→坐标回写+按中心点匹配OCR文本"""

    def __init__(self, config: Dict[str, Any], ocr_engine: Any):
        self.config = config or {}
        self.upscale_ratio: float = self.config.get("upscale_ratio", 10 / 3)
        self.need_ocr: bool = self.config.get("need_ocr", True)
        self.row_threshold: int = self.config.get("row_threshold", 10)
        self.col_threshold: int = self.config.get("col_threshold", 15)
        self.ocr_conf_threshold: float = self.config.get("ocr_conf_threshold", 0.5)
        self.cell_crop_margin: int = self.config.get("cell_crop_margin", 2)
        self.table_model = UnetTableModel(ocr_engine)
        self.ocr_engine = ocr_engine

    @staticmethod
    def _to_unet_ocr_format(ocr_boxes: List[Dict[str, Any]]) -> List[List[Any]]:
        """将OCR结果转成 UNet 期望格式 [[poly4,text,score], ...]，坐标用浮点。"""
        formatted = []
        for item in ocr_boxes:
            poly = item.get("bbox", [])
            text = item.get("text", "")
            score = item.get("confidence", 0.0)
            if not poly or len(poly) < 4:
                continue
            # 统一成4点 (4,2)
            if len(poly) == 8:
                poly_pts = [[float(poly[i]), float(poly[i + 1])] for i in range(0, 8, 2)]
            elif len(poly) == 4:
                x1, y1, x2, y2 = poly
                poly_pts = [[x1, y1], [x2, y1], [x2, y2], [x1, y2]]
            else:
                # 其他格式跳过
                continue
            formatted.append([poly_pts, text, float(score)])
        return formatted

    @staticmethod
    def _scale_poly(poly: List[List[float]], ratio: float) -> List[List[float]]:
        return [[p[0] * ratio, p[1] * ratio] for p in poly]

    @staticmethod
    def _poly_to_bbox(poly: np.ndarray) -> List[float]:
        xs = poly[:, 0]
        ys = poly[:, 1]
        return [float(xs.min()), float(ys.min()), float(xs.max()), float(ys.max())]

    def _match_text_by_center(
        self,
        cells_bbox: List[List[float]],
        ocr_boxes: List[Dict[str, Any]],
    ) -> List[str]:
        """使用中心点落格分配文本，行内按 y 排序后拼接。"""
        texts_per_cell: List[str] = []
        centers = []
        for item in ocr_boxes:
            poly = item.get("bbox", [])
            if not poly:
                continue
            if len(poly) == 8:
                xs = [poly[i] for i in range(0, 8, 2)]
                ys = [poly[i] for i in range(1, 8, 2)]
                cx = (min(xs) + max(xs)) / 2
                cy = (min(ys) + max(ys)) / 2
            elif len(poly) == 4:
                x1, y1, x2, y2 = poly
                cx = (x1 + x2) / 2
                cy = (y1 + y2) / 2
            else:
                continue
            centers.append((cx, cy, item.get("text", ""), item.get("confidence", 0.0)))

        for bbox in cells_bbox:
            x1, y1, x2, y2 = bbox
            collected = [(t, cy) for cx, cy, t, conf in centers if x1 <= cx <= x2 and y1 <= cy <= y2]
            collected.sort(key=lambda x: x[1])
            cell_text = " ".join([t for t, _ in collected]) if collected else ""
            texts_per_cell.append(cell_text)
        return texts_per_cell

    def recognize(
        self,
        table_image: np.ndarray,
        ocr_boxes: List[Dict[str, Any]],
    ) -> Dict[str, Any]:
        """运行有线表格识别并输出 HTML + cells（原图坐标）。"""
        # 预放大
        upscale = self.upscale_ratio if self.upscale_ratio and self.upscale_ratio > 0 else 1.0
        h, w = table_image.shape[:2]
        if upscale != 1.0:
            table_image_up = cv2.resize(table_image, (int(w * upscale), int(h * upscale)), interpolation=cv2.INTER_CUBIC)
        else:
            table_image_up = table_image

        # OCR 坐标缩放到放大后
        ocr_scaled = []
        for item in ocr_boxes or []:
            poly = item.get("bbox", [])
            if not poly:
                continue
            if len(poly) == 8:
                scaled = [p * upscale for p in poly]
            elif len(poly) == 4:
                x1, y1, x2, y2 = poly
                scaled = [x1 * upscale, y1 * upscale, x2 * upscale, y2 * upscale]
            else:
                continue
            ocr_scaled.append({"bbox": scaled, "text": item.get("text", ""), "confidence": item.get("confidence", 0.0)})

        # 组装 UNet OCR 输入格式
        unet_ocr = self._to_unet_ocr_format(ocr_scaled)

        # 调用 UNet - 使用单次 wired_table_model 调用获取所有结果
        html_code = ""
        cell_polys_flat = None
        logic_points = None
        try:
            # 类型忽略：UnetTableModel 期望 [[poly,text,score], ...]
            # 单次调用 wired_table_model，返回 WiredTableOutput namedtuple
            # 包含: pred_html, cell_bboxes, logic_points, elapse
            wired_out = self.table_model.wired_table_model(
                table_image_up,
                unet_ocr,  # type: ignore[arg-type]
                col_threshold=self.col_threshold,
                row_threshold=self.row_threshold,
            )
            # 统一从 wired_out 提取所有数据
            html_code = wired_out.pred_html if wired_out.pred_html else ""
            cell_polys_flat = wired_out.cell_bboxes
            logic_points = wired_out.logic_points
        except Exception:
            html_code = ""

        cells = []
        if cell_polys_flat is not None and logic_points is not None:
            polys = np.array(cell_polys_flat, dtype=float).reshape(-1, 4, 2)
            # 缩回裁剪坐标
            polys /= upscale
            bboxes = [self._poly_to_bbox(poly) for poly in polys]
            texts = self._match_text_by_center(bboxes, ocr_boxes or [])
            # 对空文本单元格触发单元格级 OCR 补充
            if self.ocr_engine is not None and any(not t for t in texts):
                crop_list = []
                crop_info = []
                h, w = table_image.shape[:2]
                margin = self.cell_crop_margin
                for idx, bbox in enumerate(bboxes):
                    if texts[idx]:
                        continue
                    x1, y1, x2, y2 = bbox
                    x1i, y1i, x2i, y2i = map(int, [x1, y1, x2, y2])
                    
                    # 增加裁剪边距防止文字被截断（特别是边界字符如"司"）
                    x1i = max(0, x1i - margin)
                    y1i = max(0, y1i - margin)
                    x2i = min(w, x2i + margin)
                    y2i = min(h, y2i + margin)
                    
                    if x2i <= x1i or y2i <= y1i:
                        continue
                    crop = table_image[y1i:y2i, x1i:x2i]
                    if crop.size == 0:
                        continue
                    crop_list.append(crop)
                    crop_info.append(idx)
                if crop_list:
                    try:
                        ocr_res = self.ocr_engine.ocr(crop_list, det=False)
                        if ocr_res and isinstance(ocr_res, list) and len(ocr_res) == 1:
                            for loc, (text, score) in zip(crop_info, ocr_res[0]):
                                if score >= self.ocr_conf_threshold and text:
                                    texts[loc] = text
                    except Exception:
                        pass
            for idx, bbox in enumerate(bboxes):
                lp = logic_points[idx] if len(logic_points) > idx else [0, 0, 0, 0]
                cells.append({
                    "bbox": bbox,
                    "row": int(lp[0]) if len(lp) > 0 else 0,
                    "col": int(lp[2]) if len(lp) > 2 else 0,
                    "text": texts[idx] if idx < len(texts) else "",
                    "matched_text": texts[idx] if idx < len(texts) else "",
                    "score": 100.0,
                })

        # 通过BeautifulSoup增强HTML，添加data-bbox和data-score属性（保留原始HTML结构）
        html_enhanced = self._enhance_html_with_cell_data(html_code, cells)

        return {
            "html": html_enhanced or html_code or "",
            "cells": cells,
        }
    
    def _enhance_html_with_cell_data(self, html_code: str, cells: List[Dict[str, Any]]) -> str:
        """通过BeautifulSoup增强HTML，为每个td添加data-bbox和data-score属性
        
        保留原始MinerU的rowspan/colspan等属性，只增加定位信息。按照cells的row/col与HTML中td的位置关系进行匹配。
        
        Args:
            html_code: MinerU生成的原始HTML
            cells: 单元格列表，包含bbox、row、col等信息
        
        Returns:
            增强后的HTML字符串，包含data-bbox和data-score属性
        """
        if not html_code or not cells:
            return html_code
        
        try:
            from bs4 import BeautifulSoup
        except ImportError:
            return html_code
        
        soup = BeautifulSoup(html_code, 'html.parser')
        
        # 建立cell快速查询字典：(row, col) -> cell
        cell_dict = {}
        for cell in cells:
            row = cell.get("row", 0)
            col = cell.get("col", 0)
            key = (row, col)
            cell_dict[key] = cell
        
        # 遍历HTML中的所有tr和td，按行列顺序进行匹配
        rows = soup.find_all('tr')
        for row_idx, tr in enumerate(rows):
            tds = tr.find_all('td')  # type: ignore
            col_idx = 0
            for td in tds:
                # 获取colspan和rowspan属性
                colspan_str = td.get('colspan')  # type: ignore
                rowspan_str = td.get('rowspan')  # type: ignore
                try:
                    colspan = int(str(colspan_str)) if colspan_str else 1
                    rowspan = int(str(rowspan_str)) if rowspan_str else 1
                except (ValueError, TypeError):
                    colspan = 1
                    rowspan = 1
                
                # 根据row_idx和col_idx查找对应的cell
                cell_key = (row_idx, col_idx)
                if cell_key in cell_dict:
                    cell = cell_dict[cell_key]
                    bbox = cell.get("bbox", [])
                    score = cell.get("score", 100.0)
                    
                    # 添加data-bbox属性
                    if bbox and len(bbox) >= 4:
                        bbox_str = ",".join(map(str, map(int, bbox[:4])))
                        td['data-bbox'] = f"[{bbox_str}]"  # type: ignore
                    
                    # 添加data-score属性
                    td['data-score'] = f"{score:.4f}"  # type: ignore
                
                # 更新列索引（考虑colspan）
                col_idx += colspan
        
        return str(soup)
