"""
PDF 结构提取器 - 异步并发 OCR 版本

基于 splitter_pdf 逻辑，直接提取章节结构并记录页码。
支持 OCR 增强：检测表格区域并使用 asyncio 5并发 OCR，其他文本保持 PyMuPDF 提取。
输出格式兼容后续分类与组装流程。
"""

import asyncio
import base64
import io
import re
from dataclasses import dataclass
from typing import Dict, Any, List, Optional, Tuple, Set

import fitz
import numpy as np
import requests

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

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


@dataclass
class TableRegion:
    """表格区域信息"""
    page_num: int
    page: fitz.Page
    bbox: Tuple[float, float, float, float]
    score: float


@dataclass
class OcrResult:
    """OCR 结果"""
    page_num: int
    bbox: Tuple[float, float, float, float]
    score: float
    text: str
    success: bool


class PdfStructureExtractor:
    """PDF 章节结构提取器（支持 OCR 异步并发）"""

    CHAPTER_PATTERN = re.compile(r"^第[一二三四五六七八九十百]+章\s*.*")
    SECTION_PATTERN = re.compile(r"^[一二三四五六七八九十百]+、\s*.*")
    TOC_PATTERN = re.compile(r"\.{3,}|…{2,}")

    # OCR 配置
    MAX_SHORT_EDGE = 1024
    JPEG_QUALITY = 90
    OCR_DPI = 200
    OCR_CONFIDENCE_THRESHOLD = 0.5
    OCR_CONCURRENT_WORKERS = 5

    def __init__(
        self,
        clip_top: float = 60,
        clip_bottom: float = 60,
        use_ocr: bool = False,
        ocr_api_url: str = "http://183.220.37.46:25429/v1/chat/completions",
        ocr_timeout: int = 600,
        ocr_api_key: str = "",
    ):
        self.clip_top = clip_top
        self.clip_bottom = clip_bottom
        self.use_ocr = use_ocr and RAPID_LAYOUT_AVAILABLE

        # OCR 配置
        self.ocr_api_url = ocr_api_url
        self.ocr_timeout = ocr_timeout
        self.ocr_api_key = ocr_api_key
        self._layout_engine: Optional[Any] = None

        if use_ocr and not RAPID_LAYOUT_AVAILABLE:
            logger.warning("RapidLayout 未安装，OCR 功能不可用")

    def _get_layout_engine(self) -> Optional[Any]:
        """延迟初始化 RapidLayout"""
        if self._layout_engine is None and RAPID_LAYOUT_AVAILABLE:
            self._layout_engine = RapidLayout()
        return self._layout_engine

    def extract(self, file_content: bytes) -> Dict[str, Any]:
        """
        从 PDF 字节流提取章节结构。

        Returns:
            {
                "chapters": {
                    "第一章 xxx": {
                        "章节标题": {"content": "...", "page_start": 1, "page_end": 1},
                        "一、xxx": {"content": "...", "page_start": 2, "page_end": 3},
                    }
                },
                "total_pages": N
            }
        """
        doc = fitz.open(stream=file_content)
        try:
            structure = self._extract_from_doc(doc)
            structure["total_pages"] = len(doc)
            return structure
        finally:
            doc.close()

    def _extract_from_doc(self, doc: fitz.Document) -> Dict[str, Any]:
        """提取文档结构（支持 OCR 异步并发）"""

        # === 阶段1: 收集所有需要 OCR 的表格区域 ===
        table_regions: List[TableRegion] = []

        if self.use_ocr:
            logger.info("[OCR预处理] 扫描所有页面的表格区域...")
            for page_num in range(len(doc)):
                page = doc.load_page(page_num)
                rect = page.rect
                clip_box = fitz.Rect(0, self.clip_top, rect.width, rect.height - self.clip_bottom)
                regions = self._detect_table_regions(page, page_num + 1, clip_box)
                for bbox, score in regions:
                    table_regions.append(TableRegion(
                        page_num=page_num + 1,
                        page=page,
                        bbox=bbox,
                        score=score
                    ))
            logger.info(f"[OCR预处理] 共发现 {len(table_regions)} 个表格区域需要 OCR")

        # === 阶段2: 异步并发执行 OCR (5并发) ===
        ocr_results: List[OcrResult] = []

        if table_regions:
            logger.info(f"[OCR执行] 使用 {self.OCR_CONCURRENT_WORKERS} 并发执行 OCR...")

            # 使用 nest_asyncio 允许嵌套事件循环
            import nest_asyncio
            nest_asyncio.apply()

            ocr_results = asyncio.run(self._process_ocr_concurrent(table_regions))

            success_count = sum(1 for r in ocr_results if r.success)
            logger.info(f"[OCR执行] 完成 {success_count}/{len(table_regions)} 个表格 OCR")

        # 按页码分组 OCR 结果
        ocr_by_page: Dict[int, List[OcrResult]] = {}
        for result in ocr_results:
            if result.success:
                if result.page_num not in ocr_by_page:
                    ocr_by_page[result.page_num] = []
                ocr_by_page[result.page_num].append(result)

        # === 阶段3: 提取页面文本（应用 OCR 结果）并切分章节 ===
        structured_data: Dict[str, Dict[str, Dict[str, Any]]] = {}
        current_chapter = "未分类前言"
        current_section = "默认部分"
        in_body = False

        logger.info("[文本提取] 提取页面内容并切分章节...")

        for page_num in range(len(doc)):
            page = doc.load_page(page_num)
            rect = page.rect
            clip_box = fitz.Rect(0, self.clip_top, rect.width, rect.height - self.clip_bottom)

            # 获取页面文本（应用 OCR 结果）
            if page_num + 1 in ocr_by_page:
                original_text = page.get_text("text", clip=clip_box)
                ocr_results_list = [
                    {
                        "region_index": i,
                        "bbox": r.bbox,
                        "score": r.score,
                        "ocr_text": r.text,
                    }
                    for i, r in enumerate(ocr_by_page[page_num + 1])
                ]
                text = self._replace_table_regions(page, original_text, ocr_results_list, clip_box)
            else:
                text = page.get_text("text", clip=clip_box)

            lines = text.split("\n")

            for line in lines:
                line = line.strip()
                if not line:
                    continue
                if self._is_header_footer(line):
                    continue

                # 跳过目录阶段
                if not in_body:
                    if self.CHAPTER_PATTERN.match(line) and not self.TOC_PATTERN.search(line):
                        in_body = True
                    else:
                        continue

                # 跳过残余目录格式
                if self.TOC_PATTERN.search(line):
                    continue

                # 匹配章标题
                if self.CHAPTER_PATTERN.match(line):
                    current_chapter = self._clean_chapter_title(line)
                    current_section = "章节标题"
                    if current_chapter not in structured_data:
                        structured_data[current_chapter] = {}
                    if current_section not in structured_data[current_chapter]:
                        structured_data[current_chapter][current_section] = {
                            "lines": [],
                            "page_start": page_num + 1,
                            "page_end": page_num + 1,
                        }
                    continue

                # 匹配节标题
                if self.SECTION_PATTERN.match(line):
                    current_section = line
                    if current_chapter not in structured_data:
                        structured_data[current_chapter] = {}
                    if current_section not in structured_data[current_chapter]:
                        structured_data[current_chapter][current_section] = {
                            "lines": [],
                            "page_start": page_num + 1,
                            "page_end": page_num + 1,
                        }
                    continue

                # 确保结构存在
                if current_chapter not in structured_data:
                    structured_data[current_chapter] = {}
                if current_section not in structured_data[current_chapter]:
                    structured_data[current_chapter][current_section] = {
                        "lines": [],
                        "page_start": page_num + 1,
                        "page_end": page_num + 1,
                    }

                # 添加内容
                structured_data[current_chapter][current_section]["lines"].append(line)
                structured_data[current_chapter][current_section]["page_end"] = page_num + 1

        # 将行列表拼接为文本
        result: Dict[str, Any] = {"chapters": {}}
        for chap, sections in structured_data.items():
            result["chapters"][chap] = {}
            for sec, data in sections.items():
                result["chapters"][chap][sec] = {
                    "content": "\n".join(data["lines"]),
                    "page_start": data["page_start"],
                    "page_end": data["page_end"],
                }

        logger.info(f"[PdfExtractor] 提取完成，共 {len(result['chapters'])} 个章节")
        return result

    async def _process_ocr_concurrent(self, regions: List[TableRegion]) -> List[OcrResult]:
        """异步并发处理 OCR"""
        semaphore = asyncio.Semaphore(self.OCR_CONCURRENT_WORKERS)

        async def _ocr_single(region: TableRegion) -> OcrResult:
            async with semaphore:
                try:
                    text = await self._ocr_region_async(region.page, region.bbox)
                    return OcrResult(
                        page_num=region.page_num,
                        bbox=region.bbox,
                        score=region.score,
                        text=text,
                        success=True,
                    )
                except Exception as e:
                    logger.error(f"  第 {region.page_num} 页表格 OCR 失败: {e}")
                    return OcrResult(
                        page_num=region.page_num,
                        bbox=region.bbox,
                        score=region.score,
                        text="",
                        success=False,
                    )

        tasks = [_ocr_single(r) for r in regions]
        return await asyncio.gather(*tasks)

    async def _ocr_region_async(
        self,
        page: fitz.Page,
        bbox: Tuple[float, float, float, float]
    ) -> str:
        """异步执行单区域 OCR"""
        loop = asyncio.get_event_loop()
        return await loop.run_in_executor(None, self._ocr_table_region, page, bbox)

    def _detect_table_regions(
        self,
        page: fitz.Page,
        page_num: int,
        clip_box: fitz.Rect
    ) -> List[Tuple[Tuple[float, float, float, float], float]]:
        """检测页面中的表格区域，返回坐标列表"""
        table_regions: List[Tuple[Tuple[float, float, float, float], float]] = []

        if not RAPID_LAYOUT_AVAILABLE:
            return table_regions

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

        # 渲染页面（裁剪区域）
        pix = page.get_pixmap(dpi=self.OCR_DPI, clip=clip_box)
        img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.height, pix.width, 3)

        try:
            layout_output = layout_engine(img)

            # 解析版面结果
            if hasattr(layout_output, 'boxes') and hasattr(layout_output, 'class_names'):
                # 获取缩放比例
                scale_x = clip_box.width / img.shape[1]
                scale_y = clip_box.height / img.shape[0]

                for box, label, score in zip(layout_output.boxes, layout_output.class_names, layout_output.scores):
                    if label == "table" and score > self.OCR_CONFIDENCE_THRESHOLD:
                        # 转换为 PDF 坐标
                        pdf_x1 = clip_box.x0 + box[0] * scale_x
                        pdf_y1 = clip_box.y0 + box[1] * scale_y
                        pdf_x2 = clip_box.x0 + box[2] * scale_x
                        pdf_y2 = clip_box.y0 + box[3] * scale_y

                        table_regions.append(((pdf_x1, pdf_y1, pdf_x2, pdf_y2), score))

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

        return table_regions

    def _ocr_table_region(self, page: fitz.Page, bbox: Tuple[float, float, float, float]) -> str:
        """对指定区域进行 OCR 识别（使用 GLM-OCR）"""
        # 渲染指定区域
        rect = fitz.Rect(bbox)
        pix = page.get_pixmap(dpi=self.OCR_DPI, clip=rect)
        img_bytes = pix.tobytes("jpeg")

        # 压缩图片
        compressed = self._compress_image(img_bytes)
        img_base64 = base64.b64encode(compressed).decode('utf-8')

        # 请求 OCR
        payload = {
            "model": "GLM-OCR",
            "messages": [
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "识别图片中的表格内容，按原文排版输出。"
                                    "注意："
                                    "1. 表格用 Markdown 表格格式"
                                    "2. 保持换行和列对齐"
                                    "3. 只输出表格内容，不要其他说明"
                        },
                        {
                            "type": "image_url",
                            "image_url": {"url": f"data:image/jpeg;base64,{img_base64}"}
                        }
                    ]
                }
            ],
            "max_tokens": 2048,
            "temperature": 0.1
        }

        headers = {"Content-Type": "application/json"}
        if self.ocr_api_key:
            headers["Authorization"] = f"Bearer {self.ocr_api_key}"

        response = requests.post(
            self.ocr_api_url,
            headers=headers,
            json=payload,
            timeout=self.ocr_timeout
        )
        response.raise_for_status()

        result = response.json()
        return self._extract_ocr_content(result)

    def _replace_table_regions(
        self,
        page: fitz.Page,
        original_text: str,
        ocr_results: List[Dict],
        clip_box: fitz.Rect
    ) -> str:
        """用 OCR 结果替换原始文本中的表格区域"""
        if not ocr_results:
            return original_text

        # 获取页面上的文本块及其坐标
        text_blocks = []
        for block in page.get_text("blocks"):
            x0, y0, x1, y1, text, _, _ = block
            # 只考虑裁剪区域内的文本
            if y0 >= clip_box.y0 and y1 <= clip_box.y1:
                text_blocks.append({
                    "bbox": (x0, y0, x1, y1),
                    "text": text.strip(),
                })

        # 按 Y 坐标排序
        text_blocks.sort(key=lambda b: (b["bbox"][1], b["bbox"][0]))

        # 找出属于表格区域的文本块
        replaced_indices: Set[int] = set()
        for ocr_result in ocr_results:
            bbox = ocr_result["bbox"]
            rx0, ry0, rx1, ry1 = bbox

            for idx, block in enumerate(text_blocks):
                if idx in replaced_indices:
                    continue
                bx0, by0, bx1, by1 = block["bbox"]

                # 检查重叠
                overlap_x = max(0, min(bx1, rx1) - max(bx0, rx0))
                overlap_y = max(0, min(by1, ry1) - max(by0, ry0))
                overlap_area = overlap_x * overlap_y
                block_area = (bx1 - bx0) * (by1 - by0)

                if block_area > 0 and overlap_area / block_area > 0.5:
                    replaced_indices.add(idx)

        # 构建新文本
        result_parts: List[str] = []
        last_idx = 0

        for ocr_result in sorted(ocr_results, key=lambda r: r["bbox"][1]):
            bbox = ocr_result["bbox"]
            rx0, ry0, rx1, ry1 = bbox

            # 找到该表格区域之前的文本
            region_start_idx = None
            for idx, block in enumerate(text_blocks):
                if idx in replaced_indices:
                    bx0, by0, bx1, by1 = block["bbox"]
                    if (bx0 >= rx0 - 5 and bx1 <= rx1 + 5 and
                        by0 >= ry0 - 5 and by1 <= ry1 + 5):
                        if region_start_idx is None:
                            region_start_idx = idx
                        last_idx = idx + 1

            if region_start_idx is not None:
                # 添加表格前的非表格文本
                for idx in range(last_idx - (last_idx - region_start_idx), region_start_idx):
                    if idx not in replaced_indices and idx < len(text_blocks):
                        result_parts.append(text_blocks[idx]["text"])
                        result_parts.append("\n")

                # 添加 OCR 结果
                result_parts.append(ocr_result["ocr_text"])
                result_parts.append("\n")

        # 添加剩余文本
        for idx in range(last_idx, len(text_blocks)):
            if idx not in replaced_indices:
                result_parts.append(text_blocks[idx]["text"])
                result_parts.append("\n")

        return "".join(result_parts).strip() or original_text

    def _compress_image(self, img_bytes: bytes) -> bytes:
        """压缩图片"""
        try:
            from PIL import Image
            img = Image.open(io.BytesIO(img_bytes))

            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_ocr_content(self, result: Dict) -> str:
        """从 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 ""

    @staticmethod
    def _is_header_footer(line: str) -> bool:
        return (
            "四川路桥建设集团股份有限公司" in line
            or "T梁运输及安装专项施工方案" in line
            or line.isdigit()
        )

    @staticmethod
    def _clean_chapter_title(line: str) -> str:
        chapter_match = re.search(r"第[一二三四五六七八九十百]+章", line)
        if not chapter_match:
            return line.strip()

        prefix = chapter_match.group(0)
        remaining = line[chapter_match.end() :].strip()
        remaining = re.sub(r"^[\.\s]+", "", remaining)
        remaining = re.sub(r"\s+\d+\s*$", "", remaining)
        remaining = re.sub(r"[\._\-]{3,}[^\u4e00-\u9fa5a-zA-Z0-9]*", "", remaining)

        if remaining:
            return f"{prefix} {remaining}"
        return prefix
