"""
混合全文提取实现 - 支持 OCR

基于页数检测表格区域：
- 使用 RapidLayout 检测表格页
- 表格页走 OCR（GLM-OCR 或 MinerU）
- 其他页走本地 PyMuPDF 提取
"""

from __future__ import annotations

import base64
import io
import time
from typing import Any, Dict, List, Optional, Set

import fitz
import numpy as np
import requests

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

from ..interfaces import DocumentSource, FullTextExtractor


# 尝试导入 RapidLayout
try:
    from rapid_layout import RapidLayout
    RAPID_LAYOUT_AVAILABLE = True
except ImportError:
    RAPID_LAYOUT_AVAILABLE = False
    RapidLayout = None


def _read_ini_config(section: str, key: str, default: Any = None) -> Any:
    """从项目根目录的 config.ini 读取配置"""
    try:
        import configparser
        from pathlib import Path

        config_path = Path(__file__).parent.parent.parent.parent.parent.parent / "config" / "config.ini"
        if not config_path.exists():
            return default

        config = configparser.ConfigParser()
        config.read(config_path, encoding="utf-8")

        if section in config and key in config[section]:
            return config[section][key]
        return default
    except Exception:
        return default


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

    - 检测含表格的页面 -> OCR 识别
    - 其他页面 -> PyMuPDF 本地提取
    """

    # GLM-OCR 图片尺寸限制
    MAX_SHORT_EDGE = 1024
    JPEG_QUALITY = 90

    def __init__(
        self,
        layout_dpi: int = 200,
        ocr_dpi: int = 200,
        jpg_quality: int = 90,
        api_url: Optional[str] = None,
        timeout: int = 600,
        clip_top: float = 60,
        clip_bottom: float = 60,
    ) -> None:
        """
        初始化

        Args:
            layout_dpi: 版面分析 DPI
            ocr_dpi: OCR DPI
            jpg_quality: JPEG 质量
            api_url: OCR API 地址
            timeout: 超时时间
            clip_top: 顶部裁剪磅数
            clip_bottom: 底部裁剪磅数
        """
        self.layout_dpi = layout_dpi
        self.ocr_dpi = ocr_dpi
        self.jpg_quality = jpg_quality
        self.clip_top = clip_top
        self.clip_bottom = clip_bottom

        # OCR 引擎配置
        raw_engine = _read_ini_config("ocr", "engine", "glm_ocr")
        self.ocr_engine = raw_engine.lower().strip() if raw_engine else "glm_ocr"

        if self.ocr_engine in ("glm_ocr", "glm-ocr", "glmocr"):
            self.ocr_engine_normalized = "glm_ocr"
        elif self.ocr_engine in ("mineru", "mineru-ocr", "mineru_ocr"):
            self.ocr_engine_normalized = "mineru"
        else:
            self.ocr_engine_normalized = "glm_ocr"

        logger.info(f"[HybridExtractor] OCR 引擎: {self.ocr_engine_normalized}")

        # GLM-OCR 配置
        self.glm_api_url = api_url or _read_ini_config(
            "ocr", "glm_ocr_api_url",
            "http://183.220.37.46:25429/v1/chat/completions"
        )
        self.glm_timeout = int(_read_ini_config("ocr", "glm_ocr_timeout", "600"))
        self.glm_api_key = _read_ini_config("ocr", "glm_ocr_api_key", "")
        self.glm_headers = {"Content-Type": "application/json"}
        if self.glm_api_key:
            self.glm_headers["Authorization"] = f"Bearer {self.glm_api_key}"

        # MinerU 配置
        self.mineru_api_url = _read_ini_config(
            "ocr", "mineru_api_url",
            "http://183.220.37.46:25428/file_parse"
        )
        self.mineru_timeout = int(_read_ini_config("ocr", "mineru_timeout", "300"))

        # 版面分析引擎
        self._layout_engine: Optional[Any] = None
        self._image_cache: Dict[int, tuple] = {}

        if not RAPID_LAYOUT_AVAILABLE:
            logger.warning("RapidLayout 未安装，表格检测不可用")

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

    def _detect_table_pages(self, doc: fitz.Document) -> Set[int]:
        """
        使用 RapidLayout 检测含表格的页码

        Returns:
            包含 table 区域的页码集合（1-based）
        """
        table_pages: Set[int] = set()

        if not RAPID_LAYOUT_AVAILABLE:
            logger.warning("RapidLayout 不可用，跳过表格检测")
            return table_pages

        layout_engine = self._get_layout_engine()
        if layout_engine is None:
            return table_pages

        total_pages = len(doc)
        self._image_cache.clear()

        logger.info(f"[版面分析] 共 {total_pages} 页，DPI={self.layout_dpi}")

        for page_num in range(1, total_pages + 1):
            page = doc[page_num - 1]

            # 裁剪页眉页脚
            rect = page.rect
            clip_box = fitz.Rect(0, self.clip_top, rect.width, rect.height - self.clip_bottom)

            # 渲染页面为图片
            pix = page.get_pixmap(dpi=self.layout_dpi, clip=clip_box)
            img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)

            try:
                layout_output = layout_engine(img)

                # 解析版面结果
                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)
                    ]

                # 检测表格
                if "table" in labels:
                    table_pages.add(page_num)
                    # 缓存图片供 OCR 使用
                    jpeg_bytes = pix.tobytes("jpeg")
                    self._image_cache[page_num] = (pix.width, pix.height, jpeg_bytes)
                    logger.debug(f"  第 {page_num} 页: 检测到表格")
                else:
                    logger.debug(f"  第 {page_num} 页: 无表格")

            except Exception as e:
                logger.warning(f"  第 {page_num} 页: 版面分析失败 ({e})")

        logger.info(f"[版面分析] 完成: {len(table_pages)}/{total_pages} 页含表格")
        return table_pages

    def extract_full_text(self, source: DocumentSource) -> List[Dict[str, Any]]:
        """
        执行混合提取

        1. 使用 RapidLayout 检测表格页
        2. 表格页 -> OCR
        3. 其他页 -> PyMuPDF 本地提取
        """
        total_start = time.time()

        # 打开文档
        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)

            # 阶段 1: 版面分析检测表格页
            logger.info("[阶段1] 版面分析检测表格页...")
            layout_start = time.time()
            table_pages = self._detect_table_pages(doc)
            layout_time = time.time() - layout_start

            # 阶段 2: 分流处理
            logger.info("[阶段2] 分流处理...")
            ocr_count = 0
            ocr_total_time = 0.0

            for i, page in enumerate(doc):
                page_num = i + 1

                if page_num in table_pages:
                    # OCR 处理
                    ocr_start = time.time()
                    try:
                        if self.ocr_engine_normalized == "glm_ocr":
                            page_text = self._ocr_with_glm(page, page_num)
                        else:
                            page_text = self._ocr_with_mineru(doc, page_num)
                        ocr_total_time += time.time() - ocr_start
                        ocr_count += 1
                        logger.debug(f"  第 {page_num} 页: OCR 完成")
                    except Exception as e:
                        logger.error(f"  第 {page_num} 页: OCR 失败 ({e})，回退到本地提取")
                        page_text = self._extract_local(page)
                else:
                    # 本地提取
                    page_text = self._extract_local(page)
                    logger.debug(f"  第 {page_num} 页: 本地提取")

                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)

        finally:
            doc.close()
            self._image_cache.clear()

        # 统计输出
        total_time = time.time() - total_start
        ocr_avg = ocr_total_time / ocr_count if ocr_count > 0 else 0
        total_chars = sum(len(p["text"]) for p in pages)

        logger.info(
            f"[提取完成] 总页数: {total_pages} | "
            f"OCR: {ocr_count} | 本地: {total_pages - ocr_count} | "
            f"总耗时: {total_time:.2f}s | "
            f"版面分析: {layout_time:.2f}s | "
            f"OCR耗时: {ocr_total_time:.2f}s | "
            f"总字符: {total_chars}"
        )

        return pages

    def _extract_local(self, page: fitz.Page) -> str:
        """本地提取页面文本"""
        rect = page.rect
        clip_box = fitz.Rect(0, self.clip_top, rect.width, rect.height - self.clip_bottom)
        return page.get_text("text", clip=clip_box)

    def _ocr_with_glm(self, page: fitz.Page, page_num: int) -> str:
        """使用 GLM-OCR 识别页面"""
        # 检查缓存
        cached = self._image_cache.get(page_num)

        if cached:
            width, height, img_bytes = cached
            logger.debug(f"  [GLM-OCR] 第 {page_num} 页使用缓存图片")
        else:
            # 重新渲染
            rect = page.rect
            clip_box = fitz.Rect(0, self.clip_top, rect.width, rect.height - self.clip_bottom)
            pix = page.get_pixmap(dpi=self.ocr_dpi, clip=clip_box)
            img_bytes = pix.tobytes("jpeg")
            logger.debug(f"  [GLM-OCR] 第 {page_num} 页重新渲染")

        # 压缩图片
        compressed = self._compress_image(img_bytes)

        # Base64 编码
        img_base64 = base64.b64encode(compressed).decode('utf-8')

        # 构建请求
        payload = {
            "model": "GLM-OCR",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "请详细识别图片中的所有文字内容，保留原始排版格式，以 Markdown 格式输出。"
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"}
                        }
                    ]
                }
            ],
            "max_tokens": 2048,
            "temperature": 0.1
        }

        # 调用 API
        response = requests.post(
            self.glm_api_url,
            headers=self.glm_headers,
            json=payload,
            timeout=self.glm_timeout
        )
        response.raise_for_status()

        # 解析结果
        result = response.json()
        content = self._extract_glm_content(result)

        # 处理 HTML
        if "<table" in content.lower():
            content = self._convert_html_tables_to_markdown(content)

        return content

    def _ocr_with_mineru(self, doc: fitz.Document, page_num: int) -> str:
        """使用 MinerU 识别页面"""
        import tempfile
        import os

        # 检查缓存
        cached = self._image_cache.get(page_num)

        try:
            if cached:
                # 使用缓存的图片
                width, height, img_bytes = cached
                files = {'files': (f"page_{page_num}.jpg", io.BytesIO(img_bytes))}
            else:
                # 提取单页为临时 PDF
                single_doc = fitz.open()
                single_doc.insert_pdf(doc, from_page=page_num-1, to_page=page_num-1)

                with tempfile.NamedTemporaryFile(suffix=".pdf", delete=False) as tmp:
                    tmp_path = tmp.name

                single_doc.save(tmp_path)
                single_doc.close()

                with open(tmp_path, 'rb') as f:
                    files = {'files': (f"page_{page_num}.pdf", f)}

            response = requests.post(
                self.mineru_api_url,
                files=files,
                timeout=self.mineru_timeout
            )

            if not cached and 'tmp_path' in dir():
                try:
                    os.remove(tmp_path)
                except:
                    pass

            if response.status_code != 200:
                raise RuntimeError(f"MinerU error: {response.status_code}")

            result = response.json()
            content = ""

            if "results" in result and isinstance(result["results"], dict):
                for file_data in result["results"].values():
                    if isinstance(file_data, dict) and "md_content" in file_data:
                        content = file_data["md_content"]
                        break

            # 处理 HTML
            if "<table" in content.lower():
                content = self._convert_html_tables_to_markdown(content)

            return content

        except Exception as e:
            logger.error(f"MinerU 识别失败: {e}")
            raise

    def _compress_image(self, img_bytes: bytes) -> bytes:
        """压缩图片"""
        try:
            from PIL import Image

            img = Image.open(io.BytesIO(img_bytes))

            # 转为 RGB
            if img.mode in ('RGBA', 'LA', 'P'):
                background = Image.new('RGB', img.size, (255, 255, 255))
                if img.mode == 'P':
                    img = img.convert('RGBA')
                if img.mode in ('RGBA', 'LA'):
                    background.paste(img, mask=img.split()[-1])
                img = background
            elif img.mode != 'RGB':
                img = img.convert('RGB')

            # 缩放
            min_edge = min(img.size)
            if min_edge > self.MAX_SHORT_EDGE:
                ratio = self.MAX_SHORT_EDGE / min_edge
                new_size = (int(img.width * ratio), int(img.height * ratio))
                img = img.resize(new_size, Image.Resampling.LANCZOS)

            # 压缩
            buffer = io.BytesIO()
            img.save(buffer, format='JPEG', quality=self.JPEG_QUALITY, optimize=True)
            return buffer.getvalue()

        except Exception as e:
            logger.warning(f"图片压缩失败，使用原图: {e}")
            return img_bytes

    def _extract_glm_content(self, result: Dict[str, Any]) -> str:
        """从 GLM-OCR 响应提取内容"""
        if "choices" in result and isinstance(result["choices"], list):
            if len(result["choices"]) > 0:
                message = result["choices"][0].get("message", {})
                return message.get("content", "")
        return ""

    def _convert_html_tables_to_markdown(self, content: str) -> str:
        """将 HTML 表格转换为 Markdown"""
        import re

        def extract_cell_text(cell_html: str) -> str:
            text = re.sub(r'<[^>]+>', '', cell_html)
            text = text.replace('&nbsp;', ' ').replace('&lt;', '<').replace('&gt;', '>')
            text = text.replace('&amp;', '&').replace('&quot;', '"').replace('&#39;', "'")
            return text.strip()

        def parse_colspan(td_html: str) -> int:
            match = re.search(r'colspan=["\']?(\d+)["\']?', td_html, re.IGNORECASE)
            return int(match.group(1)) if match else 1

        def convert_table_match(match):
            table_html = match.group(0)

            # 提取行
            tr_matches = re.findall(r'<tr[^>]*>(.*?)</tr>', table_html, re.DOTALL | re.IGNORECASE)

            parsed_rows = []
            for tr_html in tr_matches:
                cells = re.findall(r'<(t[dh])[^>]*>(.*?)</\1>', tr_html, re.DOTALL | re.IGNORECASE)
                row_data = []
                for tag, cell_content in cells:
                    text = extract_cell_text(cell_content)
                    colspan = 1
                    full_cell_match = re.search(rf'<{tag}[^>]*>', tr_html)
                    if full_cell_match:
                        colspan = parse_colspan(full_cell_match.group(0))
                    row_data.append((text, colspan))
                if row_data:
                    parsed_rows.append(row_data)

            if not parsed_rows:
                return ""

            # 计算最大列数
            max_cols = max(sum(colspan for _, colspan in row) for row in parsed_rows)

            # 生成 Markdown
            md_rows = []
            for row in parsed_rows:
                expanded = []
                for text, colspan in row:
                    expanded.append(text)
                    expanded.extend([""] * (colspan - 1))
                while len(expanded) < max_cols:
                    expanded.append("")
                md_rows.append("| " + " | ".join(expanded) + " |")

            # 添加分隔行
            if len(md_rows) > 0:
                md_rows.insert(1, "| " + " | ".join(["---"] * max_cols) + " |")

            return "\n".join(md_rows)

        return re.sub(r'<table[^>]*>.*?</table>', convert_table_match, content,
                     flags=re.DOTALL | re.IGNORECASE)
