"""
混合全文提取实现 (HybridFullTextExtractor) - 飞浆版面分析版

基于飞浆 RapidLayout 版面分析，检测 table 区域判断扫描件：
1. 第一阶段：使用飞浆 RapidLayout 对所有页面进行版面分析
2. 第二阶段：含有 table 区域的页面走 MinerU OCR，其余走本地提取
"""

from __future__ import annotations

import io
import fitz  # PyMuPDF
import os
import tempfile
import numpy as np
from typing import Any, Dict, List, Optional, Set

from foundation.observability.logger.loggering import review_logger as logger

from ..config.provider import default_config_provider
from ..interfaces import DocumentSource, FullTextExtractor
from .fulltext_extractor import PdfFullTextExtractor
from .mineru_extractor import LocalMinerUFullTextExtractor

# 尝试导入 RapidLayout，如果未安装则给出友好提示
try:
    from rapid_layout import RapidLayout
    RAPID_LAYOUT_AVAILABLE = True
except ImportError:
    RAPID_LAYOUT_AVAILABLE = False
    RapidLayout = None


class HybridFullTextExtractor(FullTextExtractor):
    """
    混合提取器：基于飞浆版面分析检测 table 区域，智能路由扫描页到 MinerU OCR。
    """

    def __init__(
        self,
        layout_dpi: int = 180,
        ocr_dpi: int = 220,
        jpg_quality: int = 90
    ) -> None:
        self._cfg = default_config_provider
        # 复用已有的提取器
        self.local_extractor = PdfFullTextExtractor()
        self.mineru_extractor = LocalMinerUFullTextExtractor()  # 使用本地 MinerU

        # 飞浆版面分析配置（保守版优化参数）
        self.layout_dpi = layout_dpi      # 版面分析 DPI：180（平衡检测精度和速度）
        self.ocr_dpi = ocr_dpi            # OCR阶段 DPI：220（表格识别甜点值）
        self.jpg_quality = jpg_quality    # JPEG质量：90（几乎无损，文件可控）
        self._layout_engine: Optional[Any] = None  # 延迟初始化

        # 外部注入的进度状态字典（由 DocumentWorkflow 设置，心跳协程读取）
        # 格式：{'current': int(0-100), 'message': str}
        # 阶段一（版面分析）：current 0→50，阶段二（OCR提取）：current 50→100
        self._progress_state: Optional[dict] = None
        
        # 检查 RapidLayout 是否可用
        if not RAPID_LAYOUT_AVAILABLE:
            raise ImportError(
                "RapidLayout 未安装。请在 doc_worker_venv 虚拟环境中运行：\n"
                "pip install rapid-layout>=0.3.0"
            )

    def _get_layout_engine(self) -> Any:
        """延迟初始化 RapidLayout 引擎"""
        if self._layout_engine is None:
            logger.debug("  [初始化] 飞浆 RapidLayout 版面分析引擎...")
            self._layout_engine = RapidLayout()
        return self._layout_engine

    def _detect_table_pages(self, doc: fitz.Document, dpi: int = 150) -> Set[int]:
        """
        使用飞浆 RapidLayout 检测所有页面，返回包含 table 区域的页码集合。
        
        Args:
            doc: PyMuPDF 文档对象
            dpi: PDF 转图片的分辨率
            
        Returns:
            包含 table 区域的页码集合 (1-based)
        """
        table_pages: Set[int] = set()
        layout_engine = self._get_layout_engine()
        total_pages = len(doc)

        logger.debug(f"  [飞浆分析] 开始版面分析，共 {total_pages} 页...")

        for page_num in range(1, total_pages + 1):
            page = doc[page_num - 1]  # PyMuPDF 使用 0-based 索引

            # 1. 将页面转换为图片
            pix = page.get_pixmap(dpi=dpi)
            img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)

            # 2. 飞浆版面分析
            try:
                layout_output = layout_engine(img)

                # 3. 解析版面结果，检查是否有 table 区域
                labels = []
                if hasattr(layout_output, 'class_names'):
                    labels = list(layout_output.class_names)
                elif hasattr(layout_output, 'boxes'):
                    # 兼容不同版本的输出格式
                    labels = [
                        label for _, label, _
                        in zip(layout_output.boxes, layout_output.class_names, layout_output.scores)
                    ]

                # 4. 判断是否包含 table
                if "table" in labels:
                    table_pages.add(page_num)
                    logger.debug(f"    第 {page_num} 页: 检测到 table 区域 -> 将走 MinerU OCR")
                else:
                    region_types = ", ".join(set(labels)) if labels else "无"
                    logger.debug(f"    第 {page_num} 页: {region_types}")

            except Exception as e:
                logger.error(f"    第 {page_num} 页: 版面分析失败 ({e})，默认不走 OCR")
                # 分析失败时，保守起见不走 OCR
                pass

            # 阶段一进度：已分析页 / 总页数 → 0% ~ 50%
            if self._progress_state is not None:
                self._progress_state['current'] = int(page_num / total_pages * 50)
                self._progress_state['message'] = f"版面分析中：已分析 {page_num}/{total_pages} 页"

        logger.debug(f"  [飞浆分析] 完成，共 {len(table_pages)} 页包含 table 区域: {sorted(table_pages)}")
        return table_pages

    def extract_full_text(self, source: DocumentSource) -> List[Dict[str, Any]]:
        """
        执行混合提取流程：
        1. 首先用飞浆 RapidLayout 检测所有页面的 table 区域
        2. 含有 table 的页面走 MinerU OCR
        3. 其他页面走本地 PyMuPDF 提取
        """
        # 1. 打开文档
        if source.content is not None:
            doc = fitz.open(stream=io.BytesIO(source.content))
            source_file = "bytes_stream"
        elif source.path is not None:
            doc = fitz.open(source.path)
            source_file = str(source.path)
        else:
            raise ValueError("DocumentSource 既没有 path 也没有 content")

        pages: List[Dict[str, Any]] = []
        current_pos = 0

        try:
            total_pages = len(doc)
            logger.debug(f"开始混合提取（飞浆版面分析 + 本地 MinerU），共 {total_pages} 页...")

            if self._progress_state is not None:
                self._progress_state['current'] = 0
                self._progress_state['message'] = f"版面分析中：已分析 0/{total_pages} 页"

            # ========== 第一阶段：飞浆版面分析，检测 table 页 ==========
            table_pages = self._detect_table_pages(doc, dpi=self.layout_dpi)

            # ========== 第二阶段：分流处理 ==========
            logger.debug(f"\n开始分流处理...")
            
            for i, page in enumerate(doc):
                page_num = i + 1
                
                # 判断是否为 table 页（即扫描件）
                if page_num in table_pages:
                    logger.debug(f"  [第 {page_num} 页] 检测到 table -> 走本地 MinerU OCR")

                    # --- 扫描件处理 (MinerU OCR) ---
                    try:
                        page_text = self._ocr_page(page, page_num, source_file)
                    except Exception as e:
                        logger.error(f"    MinerU OCR 失败，回退到本地提取: {e}")
                        raw_text = page.get_text()
                        page_text = self.local_extractor._filter_header_footer(raw_text)
                else:
                    logger.debug(f"  [第 {page_num} 页] 无 table -> 走本地 PyMuPDF 提取")
                    
                    # --- 电子版处理 (本地 PyMuPDF) ---
                    text_with_tables = self.local_extractor._extract_text_with_table_placeholders(page)
                    page_text = self.local_extractor._filter_header_footer(text_with_tables)

                # --- 组装结果 ---
                pages.append({
                    "page_num": page_num,
                    "text": page_text,
                    "start_pos": current_pos,
                    "end_pos": current_pos + len(page_text),
                    "source_file": source_file
                })
                current_pos += len(page_text)

                # 阶段二进度：已处理页 / 总页数 → 50% ~ 100%
                if self._progress_state is not None:
                    self._progress_state['current'] = 50 + int(page_num / total_pages * 50)
                    ocr_flag = "（OCR）" if page_num in table_pages else ""
                    self._progress_state['message'] = f"文档提取中：已处理 {page_num}/{total_pages} 页{ocr_flag}"

        finally:
            doc.close()

        return pages

    def _ocr_page(self, page: fitz.Page, page_num: int, original_filename: str) -> str:
        """
        将单页转为图片并调用本地 MinerU OCR。
        使用 JPEG 格式以减小文件大小，提高传输效率。
        """
        # 1. 渲染为图片（保守版优化：220 DPI 提升表格识别精度）
        pix = page.get_pixmap(dpi=self.ocr_dpi)
        
        # 2. 保存为临时 JPEG 文件（比 PNG 更小）
        tmp_path = None
        try:
            with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_file:
                tmp_path = tmp_file.name
            
            # 保存为 JPEG 格式，质量 90%，几乎无损且文件可控
            pix.save(tmp_path, "jpeg", jpg_quality=self.jpg_quality)
            
            # 检查文件是否正确生成
            if not os.path.exists(tmp_path) or os.path.getsize(tmp_path) == 0:
                logger.error(f"    [WARN] 无法创建第 {page_num} 页的临时图片")
                return ""

            # 输出文件大小信息（用于调试）
            file_size_kb = os.path.getsize(tmp_path) / 1024
            logger.debug(f"    [INFO] 第 {page_num} 页图片: {file_size_kb:.1f} KB ({pix.width}x{pix.height})")
            
            # 3. 构造一个临时的 DocumentSource
            tmp_source = DocumentSource(path=tmp_path)
            
            # 4. 调用本地 MinerU
            results = self.mineru_extractor.extract_full_text(tmp_source)
            
            if results and len(results) > 0:
                return results[0]["text"]
            return ""
            
        except Exception as e:
            logger.error(f"    [WARN] 第 {page_num} 页 OCR 失败: {e}")
            return ""
            
        finally:
            # 清理临时文件
            if tmp_path and os.path.exists(tmp_path):
                try:
                    os.remove(tmp_path)
                except:
                    pass
