"""Docling Layout Predictor 适配器 (符合 BaseLayoutDetector 规范)

基于 HuggingFace transformers 直接加载 Docling 布局模型，不依赖 docling-ibm-models 包。
使用 MinerU 环境中的 transformers 库。

支持的 Hugging Face 模型仓库：
- ds4sd/docling-layout-old
- ds4sd/docling-layout-heron
- ds4sd/docling-layout-heron-101
- ds4sd/docling-layout-egret-medium
- ds4sd/docling-layout-egret-large
- ds4sd/docling-layout-egret-xlarge
"""

import cv2
import numpy as np
import threading
from pathlib import Path
from typing import Dict, List, Union, Any, Optional
from PIL import Image

try:
    from .base import BaseLayoutDetector
except ImportError:
    from base import BaseLayoutDetector

# 全局锁，防止模型初始化时的线程问题
_model_init_lock = threading.Lock()


class DoclingLayoutDetector(BaseLayoutDetector):
    """Docling Layout Predictor 适配器
    
    直接使用 transformers 库加载 Docling 布局模型，无需安装 docling-ibm-models。
    """
    
    # Docling 原始类别定义（来自 docling-ibm-models/layoutmodel/labels.py）
    # 使用 shifted canonical（带 Background）的版本，因为 RT-DETR 模型使用这个
    DOCLING_LABELS = {
        0: 'Background',
        1: 'Caption',
        2: 'Footnote',
        3: 'Formula',
        4: 'List-item',
        5: 'Page-footer',
        6: 'Page-header',
        7: 'Picture',
        8: 'Section-header',
        9: 'Table',
        10: 'Text',
        11: 'Title',
        12: 'Document Index',
        13: 'Code',
        14: 'Checkbox-Selected',
        15: 'Checkbox-Unselected',
        16: 'Form',
        17: 'Key-Value Region',
    }
    
    # 类别映射：Docling LayoutLabels → MinerU/EnhancedDocPipeline 类别体系
    # 参考：
    # - MinerU: mineru/utils/enum_class.py (BlockType, CategoryId)
    # - Pipeline: universal_doc_parser/core/pipeline_manager_v2.py (EnhancedDocPipeline 类别定义)
    CATEGORY_MAP = {
        'Caption': 'image_caption',       # Caption -> image_caption (IMAGE_TEXT_CATEGORIES)
        'Footnote': 'page_footnote',      # Footnote -> page_footnote (TEXT_CATEGORIES)
        'Formula': 'interline_equation',  # Formula -> interline_equation (EQUATION_CATEGORIES)
        'List-item': 'text',              # List-item -> text (TEXT_CATEGORIES)
        'Page-footer': 'footer',          # Page-footer -> footer (TEXT_CATEGORIES)
        'Page-header': 'header',          # Page-header -> header (TEXT_CATEGORIES)
        'Picture': 'image_body',          # Picture -> image_body (IMAGE_BODY_CATEGORIES)
        'Section-header': 'title',        # Section-header -> title (TEXT_CATEGORIES)
        'Table': 'table_body',            # Table -> table_body (TABLE_BODY_CATEGORIES)
        'Text': 'text',                   # Text -> text (TEXT_CATEGORIES)
        'Title': 'title',                 # Title -> title (TEXT_CATEGORIES)
        'Document Index': 'text',         # Document Index -> text (TEXT_CATEGORIES)
        'Code': 'code',                   # Code -> code (CODE_CATEGORIES)
        'Checkbox-Selected': 'abandon',   # Checkbox -> abandon (DISCARD_CATEGORIES)
        'Checkbox-Unselected': 'abandon', # Checkbox -> abandon (DISCARD_CATEGORIES)
        'Form': 'abandon',                # Form -> abandon (DISCARD_CATEGORIES)
        'Key-Value Region': 'text',       # Key-Value Region -> text (TEXT_CATEGORIES)
        'Background': 'abandon',          # Background -> abandon (DISCARD_CATEGORIES)
    }
    
    def __init__(self, config: Dict[str, Any]):
        """
        初始化 Docling Layout 检测器
        
        Args:
            config: 配置字典，支持以下参数：
                - model_dir: 模型目录路径或 HuggingFace 仓库 ID
                - device: 运行设备 ('cpu', 'cuda', 'mps')
                - conf: 置信度阈值 (默认 0.3)
                - num_threads: CPU 线程数 (默认 4)
        """
        super().__init__(config)
        self.model = None
        self.image_processor = None
        self._device = None
        self._threshold = 0.3
        self._num_threads = 4
        self._model_path = None
        # RT-DETR 使用 shifted labels，label_offset = 1
        self._label_offset = 1
    
    def initialize(self):
        """初始化模型"""
        try:
            import torch
            from transformers import AutoModelForObjectDetection, RTDetrImageProcessor
            from huggingface_hub import snapshot_download
            
            model_dir = self.config.get('model_dir', 'ds4sd/docling-layout-old')
            device = self.config.get('device', 'cpu')
            self._threshold = self.config.get('conf', 0.3)
            self._num_threads = self.config.get('num_threads', 4)
            
            # 设置设备
            self._device = torch.device(device)
            if device == 'cpu':
                torch.set_num_threads(self._num_threads)
            
            # 判断是本地路径还是 HuggingFace 仓库
            model_path = Path(model_dir)
            if model_path.exists() and model_path.is_dir():
                # 本地路径
                self._model_path = str(model_path)
                print(f"📂 Loading model from local path: {self._model_path}")
            else:
                # 从 HuggingFace 下载
                print(f"📥 Downloading model from HuggingFace: {model_dir}")
                self._model_path = snapshot_download(repo_id=model_dir)
            
            # 检查必要文件
            processor_config = Path(self._model_path) / "preprocessor_config.json"
            model_config = Path(self._model_path) / "config.json"
            safetensors_file = Path(self._model_path) / "model.safetensors"
            
            if not processor_config.exists():
                raise FileNotFoundError(f"Missing preprocessor_config.json in {self._model_path}")
            if not model_config.exists():
                raise FileNotFoundError(f"Missing config.json in {self._model_path}")
            if not safetensors_file.exists():
                raise FileNotFoundError(f"Missing model.safetensors in {self._model_path}")
            
            # 加载图像处理器
            self.image_processor = RTDetrImageProcessor.from_json_file(str(processor_config))
            
            # 加载模型（使用锁防止线程问题）
            with _model_init_lock:
                self.model = AutoModelForObjectDetection.from_pretrained(
                    self._model_path,
                    config=str(model_config),
                    device_map=self._device
                )
                self.model.eval()
            
            # 检测模型类型
            model_name = type(self.model).__name__
            print(f"✅ Docling Layout Detector initialized")
            print(f"   - Model: {model_name}")
            print(f"   - Device: {self._device}")
            print(f"   - Threshold: {self._threshold}")
            print(f"   - Image size: {self.image_processor.size}")
            
        except ImportError as e:
            print(f"❌ Failed to import required libraries: {e}")
            print("   Please ensure transformers and torch are installed")
            raise
        except Exception as e:
            print(f"❌ Failed to initialize Docling Layout Detector: {e}")
            raise
    
    def cleanup(self):
        """清理资源"""
        self.model = None
        self.image_processor = None
        self._model_path = None
    
    def detect(self, image: Union[np.ndarray, Image.Image]) -> List[Dict[str, Any]]:
        """
        检测布局
        
        Args:
            image: 输入图像 (numpy数组或PIL图像)
            
        Returns:
            检测结果列表，每个元素包含:
            - category: MinerU类别名称
            - bbox: [x1, y1, x2, y2]
            - confidence: 置信度
            - raw: 原始检测结果
        """
        import torch
        
        if self.model is None:
            raise RuntimeError("Model not initialized. Call initialize() first.")
        
        # 转换为 PIL Image
        if isinstance(image, np.ndarray):
            # OpenCV BGR -> RGB
            if len(image.shape) == 3 and image.shape[2] == 3:
                image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            else:
                image_rgb = image
            pil_image = Image.fromarray(image_rgb).convert("RGB")
            orig_h, orig_w = image.shape[:2]
        else:
            pil_image = image.convert("RGB")
            orig_w, orig_h = image.size
        
        # 推理
        with torch.inference_mode():
            target_sizes = torch.tensor([pil_image.size[::-1]])  # (h, w)
            inputs = self.image_processor(images=[pil_image], return_tensors="pt").to(self._device)
            outputs = self.model(**inputs)
            
            results = self.image_processor.post_process_object_detection(
                outputs,
                target_sizes=target_sizes,
                threshold=self._threshold,
            )
        
        # 解析结果
        w, h = pil_image.size
        result = results[0]
        
        formatted_results = []
        for score, label_id, box in zip(result["scores"], result["labels"], result["boxes"]):
            score = float(score.item())
            
            # 获取原始标签（考虑 offset）
            label_id_int = int(label_id.item()) + self._label_offset
            original_label = self.DOCLING_LABELS.get(label_id_int, f'unknown_{label_id_int}')
            
            # 映射到 MinerU 类别
            mineru_category = self.CATEGORY_MAP.get(original_label, 'text')
            
            # 跳过 Background
            if original_label == 'Background':
                continue
            
            # 提取边界框
            bbox_float = [float(b.item()) for b in box]
            x1 = min(w, max(0, bbox_float[0]))
            y1 = min(h, max(0, bbox_float[1]))
            x2 = min(w, max(0, bbox_float[2]))
            y2 = min(h, max(0, bbox_float[3]))
            
            bbox = [int(x1), int(y1), int(x2), int(y2)]
            
            # 计算宽高
            width = bbox[2] - bbox[0]
            height = bbox[3] - bbox[1]
            
            # 过滤太小的框
            if width < 10 or height < 10:
                continue
            
            # 过滤面积异常大的框
            area = width * height
            img_area = orig_w * orig_h
            if area > img_area * 0.95:
                continue
            
            # 生成多边形坐标
            poly = [
                bbox[0], bbox[1],  # 左上
                bbox[2], bbox[1],  # 右上
                bbox[2], bbox[3],  # 右下
                bbox[0], bbox[3],  # 左下
            ]
            
            formatted_results.append({
                'category': mineru_category,
                'bbox': bbox,
                'confidence': score,
                'raw': {
                    'original_label': original_label,
                    'original_label_id': label_id_int,
                    'poly': poly,
                    'width': width,
                    'height': height
                }
            })
        
        return formatted_results
    
    def detect_batch(
        self, 
        images: List[Union[np.ndarray, Image.Image]]
    ) -> List[List[Dict[str, Any]]]:
        """
        批量检测布局（更高效）
        
        Args:
            images: 输入图像列表
            
        Returns:
            每个图像的检测结果列表
        """
        import torch
        
        if self.model is None:
            raise RuntimeError("Model not initialized. Call initialize() first.")
        
        if not images:
            return []
        
        # 转换为 PIL Image 列表
        pil_images = []
        orig_sizes = []
        
        for image in images:
            if isinstance(image, np.ndarray):
                if len(image.shape) == 3 and image.shape[2] == 3:
                    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                else:
                    image_rgb = image
                pil_images.append(Image.fromarray(image_rgb).convert("RGB"))
                orig_sizes.append((image.shape[1], image.shape[0]))  # (w, h)
            else:
                pil_images.append(image.convert("RGB"))
                orig_sizes.append(image.size)  # (w, h)
        
        # 批量推理
        with torch.inference_mode():
            target_sizes = torch.tensor([img.size[::-1] for img in pil_images])
            inputs = self.image_processor(images=pil_images, return_tensors="pt").to(self._device)
            outputs = self.model(**inputs)
            
            results_list = self.image_processor.post_process_object_detection(
                outputs,
                target_sizes=target_sizes,
                threshold=self._threshold,
            )
        
        # 转换结果
        all_formatted_results = []
        
        for pil_img, results, (orig_w, orig_h) in zip(pil_images, results_list, orig_sizes):
            w, h = pil_img.size
            formatted_results = []
            
            for score, label_id, box in zip(results["scores"], results["labels"], results["boxes"]):
                score = float(score.item())
                
                label_id_int = int(label_id.item()) + self._label_offset
                original_label = self.DOCLING_LABELS.get(label_id_int, f'unknown_{label_id_int}')
                mineru_category = self.CATEGORY_MAP.get(original_label, 'text')
                
                if original_label == 'Background':
                    continue
                
                bbox_float = [float(b.item()) for b in box]
                x1 = min(w, max(0, bbox_float[0]))
                y1 = min(h, max(0, bbox_float[1]))
                x2 = min(w, max(0, bbox_float[2]))
                y2 = min(h, max(0, bbox_float[3]))
                
                bbox = [int(x1), int(y1), int(x2), int(y2)]
                width = bbox[2] - bbox[0]
                height = bbox[3] - bbox[1]
                
                if width < 10 or height < 10:
                    continue
                
                area = width * height
                img_area = orig_w * orig_h
                if area > img_area * 0.95:
                    continue
                
                poly = [
                    bbox[0], bbox[1],
                    bbox[2], bbox[1],
                    bbox[2], bbox[3],
                    bbox[0], bbox[3],
                ]
                
                formatted_results.append({
                    'category': mineru_category,
                    'bbox': bbox,
                    'confidence': score,
                    'raw': {
                        'original_label': original_label,
                        'original_label_id': label_id_int,
                        'poly': poly,
                        'width': width,
                        'height': height
                    }
                })
            
            all_formatted_results.append(formatted_results)
        
        return all_formatted_results
    
    def visualize(
        self, 
        img: np.ndarray, 
        results: List[Dict],
        output_path: str = None,
        show_confidence: bool = True,
        min_confidence: float = 0.0
    ) -> np.ndarray:
        """
        可视化检测结果
        
        Args:
            img: 输入图像 (BGR 格式)
            results: 检测结果 (MinerU 格式)
            output_path: 输出路径（可选）
            show_confidence: 是否显示置信度
            min_confidence: 最小置信度阈值
            
        Returns:
            标注后的图像
        """
        import random
        
        vis_img = img.copy()
        
        # 预定义类别颜色（与 EnhancedDocPipeline 保持一致）
        predefined_colors = {
            # 文本类
            'text': (153, 0, 76),
            'title': (102, 102, 255),
            'header': (128, 128, 128),
            'footer': (128, 128, 128),
            'page_footnote': (200, 200, 200),
            # 表格类
            'table_body': (204, 204, 0),
            'table_caption': (255, 255, 102),
            # 图片类
            'image_body': (153, 255, 51),
            'image_caption': (102, 178, 255),
            # 公式类
            'interline_equation': (0, 255, 0),
            # 代码类
            'code': (102, 0, 204),
            # 丢弃类
            'abandon': (100, 100, 100),
        }
        
        # 过滤低置信度结果
        filtered_results = [
            res for res in results 
            if res['confidence'] >= min_confidence
        ]
        
        if not filtered_results:
            print(f"⚠️ No results to visualize (min_confidence={min_confidence})")
            return vis_img
        
        # 为每个出现的类别分配颜色
        category_colors = {}
        for res in filtered_results:
            cat = res['category']
            if cat not in category_colors:
                if cat in predefined_colors:
                    category_colors[cat] = predefined_colors[cat]
                else:
                    category_colors[cat] = (
                        random.randint(50, 255),
                        random.randint(50, 255),
                        random.randint(50, 255)
                    )
        
        # 绘制检测框
        for res in filtered_results:
            bbox = res['bbox']
            x1, y1, x2, y2 = bbox
            cat = res['category']
            confidence = res['confidence']
            color = category_colors[cat]
            
            # 获取原始标签
            original_label = res.get('raw', {}).get('original_label', cat)
            
            # 绘制矩形边框
            cv2.rectangle(vis_img, (x1, y1), (x2, y2), color, 2)
            
            # 构造标签文本
            if show_confidence:
                label = f"{original_label}->{cat} {confidence:.2f}"
            else:
                label = f"{original_label}->{cat}"
            
            # 计算标签尺寸
            label_size, baseline = cv2.getTextSize(
                label, 
                cv2.FONT_HERSHEY_SIMPLEX, 
                0.4, 
                1
            )
            label_w, label_h = label_size
            
            # 绘制标签背景
            cv2.rectangle(
                vis_img,
                (x1, y1 - label_h - 4),
                (x1 + label_w, y1),
                color,
                -1
            )
            
            # 绘制标签文字
            cv2.putText(
                vis_img,
                label,
                (x1, y1 - 2),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4,
                (255, 255, 255),
                1,
                cv2.LINE_AA
            )
        
        # 添加图例
        if category_colors:
            self._draw_legend(vis_img, category_colors, len(filtered_results))
        
        # 保存可视化结果
        if output_path:
            output_path = Path(output_path)
            output_path.parent.mkdir(parents=True, exist_ok=True)
            cv2.imwrite(str(output_path), vis_img)
            print(f"💾 Visualization saved to: {output_path}")
        
        return vis_img
    
    def _draw_legend(
        self, 
        img: np.ndarray, 
        category_colors: Dict[str, tuple],
        total_count: int
    ):
        """在图像上绘制图例"""
        legend_x = img.shape[1] - 200
        legend_y = 20
        line_height = 25
        
        # 绘制半透明背景
        overlay = img.copy()
        cv2.rectangle(
            overlay,
            (legend_x - 10, legend_y - 10),
            (img.shape[1] - 10, legend_y + len(category_colors) * line_height + 30),
            (255, 255, 255),
            -1
        )
        cv2.addWeighted(overlay, 0.7, img, 0.3, 0, img)
        
        # 绘制标题
        cv2.putText(
            img,
            f"Legend ({total_count} total)",
            (legend_x, legend_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.5,
            (0, 0, 0),
            1,
            cv2.LINE_AA
        )
        
        # 绘制每个类别
        y_offset = legend_y + line_height
        for cat, color in sorted(category_colors.items()):
            cv2.rectangle(
                img,
                (legend_x, y_offset - 10),
                (legend_x + 15, y_offset),
                color,
                -1
            )
            cv2.rectangle(
                img,
                (legend_x, y_offset - 10),
                (legend_x + 15, y_offset),
                (0, 0, 0),
                1
            )
            
            cv2.putText(
                img,
                cat,
                (legend_x + 20, y_offset - 2),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.4,
                (0, 0, 0),
                1,
                cv2.LINE_AA
            )
            
            y_offset += line_height


# 测试代码
if __name__ == "__main__":
    import sys
    
    # 测试配置 - 使用 HuggingFace 下载模型
    config = {
        'model_dir': 'ds4sd/docling-layout-old',  # HuggingFace 仓库 ID
        'device': 'cpu',
        'conf': 0.3,
        'num_threads': 4
    }
    
    # 初始化检测器
    print("🔧 Initializing Docling Layout Detector...")
    detector = DoclingLayoutDetector(config)
    detector.initialize()
    
    # 读取测试图像
    img_path = "/Users/zhch158/workspace/data/流水分析/康强_北京农村商业银行/ppstructurev3_client_results/康强_北京农村商业银行/康强_北京农村商业银行_page_001.png"
    
    print(f"\n📖 Loading image: {img_path}")
    img = cv2.imread(img_path)
    
    if img is None:
        print(f"❌ Failed to load image: {img_path}")
        sys.exit(1)
    
    print(f"   Image shape: {img.shape}")
    
    # 执行检测
    print("\n🔍 Detecting layout...")
    results = detector.detect(img)
    
    print(f"\n✅ 检测到 {len(results)} 个区域:")
    for i, res in enumerate(results, 1):
        print(f"  [{i}] {res['category']}: "
              f"score={res['confidence']:.3f}, "
              f"bbox={res['bbox']}, "
              f"original={res['raw']['original_label']}")
    
    # 统计各类别
    category_counts = {}
    for res in results:
        cat = res['category']
        category_counts[cat] = category_counts.get(cat, 0) + 1
    
    print(f"\n📊 类别统计 (MinerU格式):")
    for cat, count in sorted(category_counts.items()):
        print(f"  - {cat}: {count}")
    
    # 可视化
    if len(results) > 0:
        print("\n🎨 Generating visualization...")
        
        output_dir = Path(__file__).parent.parent.parent / "tests" / "output"
        output_dir.mkdir(parents=True, exist_ok=True)
        output_path = output_dir / f"{Path(img_path).stem}_docling_layout_vis.jpg"
        
        vis_img = detector.visualize(
            img, 
            results, 
            output_path=str(output_path),
            show_confidence=True,
            min_confidence=0.0
        )
        
        print(f"💾 Visualization saved to: {output_path}")
    
    # 清理
    detector.cleanup()
    print("\n✅ 测试完成!")
