"""PDF 图片与表格提取 — 基于 caption 定位的页面区域截图。

核心思路：学术论文排版极其规整，Figure caption 在图下方，Table caption 在表格上方。
因此反过来：先找 caption 文字 → 向上/向下截取页面区域 → 渲染为 PNG。

优势（相比提取嵌入位图）：
- 复合图表不会被拆成碎片（整块截取）
- 矢量图也能截取（页面渲染包含一切）
- 不依赖 find_tables()（纯文本匹配 caption）
"""

from __future__ import annotations

import json
import logging
import re
from pathlib import Path

from app.services.pdf_downloader import paper_dir
from app.utils import TMP_DIR

logger = logging.getLogger(__name__)

# ── 截取区域参数 ───────────────────────────────────────────────────────

# Figure: caption 上方搜索图的范围（点）
_FIGURE_MAX_HEIGHT = 450  # 最大向上搜索范围
_FIGURE_MIN_HEIGHT = 50  # 最小有效截图高度
_FIGURE_DEFAULT_HEIGHT = 280  # 上方未找到内容块时的默认图高度

# Table: caption 下方搜索表格的范围
_TABLE_MAX_HEIGHT = 500  # 最大向下搜索范围
_TABLE_MIN_HEIGHT = 30

# caption 左右扩展（双栏论文中 caption 可能比表格窄）
_REGION_SIDE_PADDING = 10
# 表格通常比 caption 文字宽，使用更大的水平扩展
_TABLE_SIDE_PADDING = 60

# 正文行距的 ~1.5 倍 ≈ 空白间隙阈值（学术论文紧密排版，30pt 太宽松）
_CONTENT_GAP_THRESHOLD = 20


# ── Caption 正则 ───────────────────────────────────────────────────────

# 要求以 Figure/Table 开头（避免匹配正文中的 "see Figure 3" 等）
# 支持三种 caption 格式：
#   "Figure 1: Title" / "Figure 1. Title" / "Figure 1 Title"（无标点，空格分隔）
# 第三种需要后续紧跟大写字母（排除 "Figure 1 shows..." 等正文引用）
_CAPTION_RE = re.compile(
    r"^(?:Fig\.?|Figure)\s+(\d+)\s*(?:[:\.]\s*|\s+(?=[A-Z]))",
    re.IGNORECASE,
)
_TABLE_CAPTION_RE = re.compile(
    r"^Table\s+(\d+)\s*(?:[:\.]\s*|\s+(?=[A-Z]))",
    re.IGNORECASE,
)

# ── 停止信号：表格边界检测遇到以下内容时立即停止 ──

# 下一个 Figure/Table caption（如 "Table 2:" "Figure 3:" "Figure 4 Title"）
_CAPTION_STOP_RE = re.compile(
    r"^(?:Table|Fig\.?|Figure)\s+\d+\s*(?:[:\.]\s*|\s+[A-Z])",
    re.IGNORECASE,
)
# Section header（如 "6.2 Evolution" "D.1 Dependency" "7 Conclusion"）
_SECTION_STOP_RE = re.compile(
    r"^(\d{1,2}(?:\.\d+)?\s+[A-Z][a-z]|[A-Z]\.\d+\s+[A-Z][a-z])"
)


def _estimate_column_x(caption: dict) -> tuple[float, float]:
    """估计 caption 所在列的水平边界（col_x0, col_x1）。

    双栏论文中 caption 宽度远小于页面宽度，据此判断左右列。
    单栏或跨栏 caption（宽度 >65% 页宽）返回整页宽度。
    caption 居中对齐（中心接近页面中线）时按跨栏处理，使用宽范围。
    """
    pw = caption["page_width"]
    caption_w = caption["caption_x1"] - caption["caption_x0"]

    # caption 宽度 >65% 页宽 → 单栏或跨栏
    if caption_w > pw * 0.65:
        return 0, pw

    cx = (caption["caption_x0"] + caption["caption_x1"]) / 2

    # caption 居中（中心距页面中线 <8%）→ 可能是跨栏表格，使用宽范围
    if abs(cx - pw / 2) / pw < 0.08:
        return (
            max(0, caption["caption_x0"] - _TABLE_SIDE_PADDING * 2),
            min(pw, caption["caption_x1"] + _TABLE_SIDE_PADDING * 2),
        )

    if cx < pw / 2:
        return 0, pw / 2
    else:
        return pw / 2, pw


def _find_captions(doc) -> list[dict]:
    """扫描整个文档，找到所有 Figure/Table caption 的位置和信息。"""
    captions = []

    for page_num in range(len(doc)):
        page = doc[page_num]
        page_width = page.rect.width
        page_height = page.rect.height
        blocks = page.get_text("blocks")

        for block in blocks:
            if len(block) < 5:
                continue
            text = str(block[4]).strip()
            if not text:
                continue

            bx0, by0, bx1, by1 = block[0], block[1], block[2], block[3]
            # 只取 block 第一行做匹配（避免 block 包含多段文字干扰）
            first_line = text.split("\n")[0].strip()

            m = _CAPTION_RE.match(first_line)
            if m:
                captions.append(
                    {
                        "type": "figure",
                        "num": int(m.group(1)),
                        "label": f"Figure {m.group(1)}",
                        "page_num": page_num,
                        "caption_y0": by0,
                        "caption_y1": by1,
                        "caption_x0": bx0,
                        "caption_x1": bx1,
                        "caption_text": text,
                        "page_width": page_width,
                        "page_height": page_height,
                    }
                )
                continue

            m = _TABLE_CAPTION_RE.match(first_line)
            if m:
                captions.append(
                    {
                        "type": "table",
                        "num": int(m.group(1)),
                        "label": f"Table {m.group(1)}",
                        "page_num": page_num,
                        "caption_y0": by0,
                        "caption_y1": by1,
                        "caption_x0": bx0,
                        "caption_x1": bx1,
                        "caption_text": text,
                        "page_width": page_width,
                        "page_height": page_height,
                    }
                )

    return captions


def _find_figure_top(page, caption: dict) -> float:
    """向上扫描页面，找到 Figure 的上边界。

    策略：
    1. 优先用嵌入图片定位（绝大多数 figure 包含嵌入图片，图片边界即 figure 边界）
    2. 无图片时回退到文本块间隙检测（处理纯矢量图如 TikZ/matplotlib PDF）
    """
    caption_y = caption["caption_y0"]
    col_x0, col_x1 = _estimate_column_x(caption)
    cx0 = max(col_x0, caption["caption_x0"] - _REGION_SIDE_PADDING)
    cx1 = min(col_x1, caption["caption_x1"] + _REGION_SIDE_PADDING)

    # 同页上方最近的 Figure/Table caption（多 figure 同页时截断）
    _caption_cutoff: float | None = None
    for b in page.get_text("blocks"):
        if len(b) < 5:
            continue
        by0, by1 = b[1], b[3]
        if by1 >= caption_y or by1 <= caption_y - _FIGURE_MAX_HEIGHT:
            continue
        first_line = str(b[4]).strip().split("\n")[0].strip()
        if _CAPTION_STOP_RE.match(first_line):
            _caption_cutoff = by0
            break

    # ── 策略 1：嵌入图片定位（覆盖绝大多数 figure） ──
    topmost_image_y: float | None = None
    for img_info in page.get_image_info():
        bbox = img_info.get("bbox")
        if bbox is None:
            continue
        if hasattr(bbox, "x0"):
            ix0, iy0, ix1, iy1 = bbox.x0, bbox.y0, bbox.x1, bbox.y1
        else:
            ix0, iy0, ix1, iy1 = bbox[0], bbox[1], bbox[2], bbox[3]
        if iy1 <= caption_y and iy1 > caption_y - _FIGURE_MAX_HEIGHT:
            if ix1 > cx0 and ix0 < cx1:
                if _caption_cutoff is not None and iy0 < _caption_cutoff:
                    continue  # 属于上方另一个 figure
                if topmost_image_y is None or iy0 < topmost_image_y:
                    topmost_image_y = iy0

    if topmost_image_y is not None:
        figure_top = topmost_image_y
    else:
        # ── 策略 2：文本块间隙检测（纯矢量图） ──
        above_blocks: list[tuple[float, float, float, float]] = []
        for b in page.get_text("blocks"):
            if len(b) < 5:
                continue
            bx0, by0, bx1, by1 = b[0], b[1], b[2], b[3]
            if by1 <= caption_y and by1 > caption_y - _FIGURE_MAX_HEIGHT:
                if bx1 > cx0 and bx0 < cx1:
                    if col_x0 > 0 and bx0 < col_x0 - _REGION_SIDE_PADDING * 2:
                        continue
                    above_blocks.append((bx0, by0, bx1, by1))

        if not above_blocks:
            return max(0, caption_y - _FIGURE_DEFAULT_HEIGHT)

        above_blocks.sort(key=lambda b: b[1], reverse=True)
        prev_bottom = caption_y
        for b in above_blocks:
            if prev_bottom - b[3] > _CONTENT_GAP_THRESHOLD:
                figure_top = prev_bottom - 5
                break
            prev_bottom = b[1]
        else:
            figure_top = above_blocks[-1][1]

    # 同页 caption 截断
    if _caption_cutoff is not None:
        figure_top = max(figure_top, _caption_cutoff + 5)

    # 限制最大高度
    if caption_y - figure_top > _FIGURE_MAX_HEIGHT:
        figure_top = caption_y - _FIGURE_MAX_HEIGHT

    return max(0, figure_top)


def _find_table_region(page, caption: dict) -> tuple[float, float, float, float]:
    """向下扫描页面，找到 Table 的下边界和水平范围。

    返回: (x0, bottom, x1) — 裁剪区域的左、下、右边界。
    上边界由调用方根据 caption 位置确定。

    策略：
    1. 收集 caption 下方的文本块（表格内容是文本）
    2. 找到连续内容区域的底部（遇到大间隙时停止）
    3. 同时检测表格内容的水平范围（表格通常比 caption 宽）
    """
    blocks = page.get_text("blocks")
    caption_y = caption["caption_y1"]  # caption 底部作为扫描起点
    caption_x0 = caption["caption_x0"]
    caption_x1 = caption["caption_x1"]
    page_height = caption["page_height"]
    page_width = caption["page_width"]

    # 估计 caption 所在列的水平边界，避免双栏论文跨列抓取
    col_x0, col_x1 = _estimate_column_x(caption)
    search_x0 = max(col_x0, caption_x0 - _TABLE_SIDE_PADDING)
    search_x1 = min(col_x1, caption_x1 + _TABLE_SIDE_PADDING)

    below_blocks: list[tuple[float, float, float, float]] = []
    for b in blocks:
        if len(b) < 5:
            continue
        bx0, by0, bx1, by1 = b[0], b[1], b[2], b[3]
        if by0 > caption_y and by0 < caption_y + _TABLE_MAX_HEIGHT:
            if bx1 > search_x0 and bx0 < search_x1:
                # 双栏论文：排除跨列正文段落（宽度 >> 列宽，起点在另一列）
                # 表格行起点在列内或列边界附近；正文段落起点在另一列（bx0 远小于 col_x0）
                if col_x0 > 0 and bx0 < col_x0 - _TABLE_SIDE_PADDING:
                    continue
                # 停止信号：遇到下一个 caption 或 section header 立即停止
                text = str(b[4]).strip()
                first_line = text.split("\n")[0].strip()
                if _CAPTION_STOP_RE.match(first_line) or _SECTION_STOP_RE.match(
                    first_line
                ):
                    break
                below_blocks.append((bx0, by0, bx1, by1))

    if not below_blocks:
        # 没有内容 → 使用默认高度和 caption 宽度
        return (
            max(0, caption_x0 - _REGION_SIDE_PADDING),
            min(page_height, caption_y + _TABLE_MIN_HEIGHT),
            min(page_width, caption_x1 + _REGION_SIDE_PADDING),
        )

    # ── 找到连续内容区域的底部 ──
    below_blocks.sort(key=lambda b: b[1])  # 按 y 升序

    prev_y = caption_y
    bottom = below_blocks[-1][3] + 5  # 最后一块的底部 + margin

    for b in below_blocks:
        gap = b[1] - prev_y  # b[1] = by0
        if gap > _CONTENT_GAP_THRESHOLD:
            bottom = prev_y + 5
            break
        prev_y = b[3]  # b[3] = by1

    # 限制最大高度
    if bottom - caption_y > _TABLE_MAX_HEIGHT:
        bottom = caption_y + _TABLE_MAX_HEIGHT

    # ── 检测表格内容的水平范围 ──
    # 只用 gap 之前的 block 计算水平范围（gap 之后的 block 属于正文，可能更宽）
    table_blocks = [b for b in below_blocks if b[1] < bottom]
    if not table_blocks:
        table_blocks = below_blocks[:1]  # 至少用第一个 block
    content_x0 = min(caption_x0, min(b[0] for b in table_blocks))
    content_x1 = max(caption_x1, max(b[2] for b in table_blocks))

    # 添加边距，不超出页面
    # 使用较小 padding，避免将相邻列内容（如同页另一列的 Figure）带入截图；
    # 同时不限制列边界 — 双栏论文中 caption 可能跨列起始
    x0 = max(0, content_x0 - _REGION_SIDE_PADDING)
    x1 = min(page_width, content_x1 + _REGION_SIDE_PADDING)

    return (x0, bottom, x1)


def extract_images_from_pdf(arxiv_id: str, pdf_path: Path | None = None) -> int:
    """从 PDF 提取 Figure/Table 截图，生成 manifest。

    策略：找 caption → 定位区域 → 渲染页面截图。

    Args:
        arxiv_id: 论文 ID
        pdf_path: PDF 路径，默认 data/tmp/{arxiv_id}/paper.pdf

    Returns:
        提取的图片数量
    """
    import pymupdf

    if pdf_path is None:
        pdf_path = TMP_DIR / arxiv_id / "paper.pdf"

    if not pdf_path.exists():
        logger.warning("PDF not found for %s: %s", arxiv_id, pdf_path)
        return 0

    images_dest = paper_dir(arxiv_id) / "images"
    images_dest.mkdir(parents=True, exist_ok=True)

    # 清理上次提取的旧图片，避免残留
    for old_file in images_dest.glob("*.png"):
        old_file.unlink()
    if (images_dest / "manifest.json").exists():
        (images_dest / "manifest.json").unlink()

    doc = pymupdf.open(str(pdf_path))
    captions = _find_captions(doc)

    if not captions:
        logger.info("No Figure/Table captions found in PDF for %s", arxiv_id)
        doc.close()
        return 0

    # 去重：同一页同一 label 可能匹配到多个 block（如正文引用 "Figure 7"）
    # 保留每个 (type, num) 的第一个匹配（即真正的 caption）
    seen_labels: dict[str, dict] = {}
    for cap in captions:
        key = cap["label"]
        if key not in seen_labels:
            seen_labels[key] = cap

    unique_captions = list(seen_labels.values())
    extracted = 0
    manifest: dict[str, dict] = {}

    zoom = 3  # 3x 渲染，保证清晰度

    for cap in unique_captions:
        page = doc[cap["page_num"]]
        pw = cap["page_width"]

        if cap["type"] == "figure":
            # Figure: caption 上方是图 → 向上找图的上边界
            top = _find_figure_top(page, cap)
            # 上方多留 5pt 边距，确保图框边框、装饰线等不被截断
            top = max(0, top - 5)
            bottom = cap["caption_y1"] + 5  # 包含 caption
            # 水平范围：caption 宽度 + 边距（图和 caption 通常等宽）
            # 但也要考虑图内容的实际宽度
            x0 = max(0, cap["caption_x0"] - _REGION_SIDE_PADDING)
            x1 = min(pw, cap["caption_x1"] + _REGION_SIDE_PADDING)

            height = bottom - top
            if height < _FIGURE_MIN_HEIGHT:
                logger.debug(
                    "Figure %s too small (%.0fpt), skipping", cap["label"], height
                )
                continue

        else:
            # Table: caption 下方是表格 → 向下找表格的下边界和水平范围
            x0, bottom, x1 = _find_table_region(page, cap)
            top = max(0, cap["caption_y0"] - 3)  # 包含 caption，上边留少许 margin

            height = bottom - top
            if height < _TABLE_MIN_HEIGHT:
                logger.debug(
                    "Table %s too small (%.0fpt), skipping", cap["label"], height
                )
                continue

        # 渲染截取
        clip = pymupdf.Rect(x0, top, x1, bottom)
        mat = pymupdf.Matrix(zoom, zoom)
        try:
            pix = page.get_pixmap(matrix=mat, clip=clip)
        except Exception:
            logger.debug("Failed to render %s region for %s", cap["label"], arxiv_id)
            continue

        filename = f"{cap['label'].replace(' ', '_').lower()}.png"
        pix.save(str(images_dest / filename))
        extracted += 1

        cap_preview = cap["caption_text"][:200] if cap["caption_text"] else ""
        manifest[filename] = {
            "page": cap["page_num"] + 1,
            "type": cap["type"],
            "label": cap["label"],
            "caption_text": cap_preview,
            "figures" if cap["type"] == "figure" else "tables": [cap["label"]],
        }
        logger.debug(
            "Rendered %s: page %d, region (%.0f,%.0f)-(%.0f,%.0f) h=%.0fpt → %s",
            cap["label"],
            cap["page_num"] + 1,
            x0,
            top,
            x1,
            bottom,
            height,
            filename,
        )

    doc.close()

    # 保存 manifest
    manifest_path = images_dest / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2))

    if extracted > 0:
        logger.info(
            "Extracted %d figure/table screenshots from PDF for %s "
            "(from %d captions found, %d unique)",
            extracted,
            arxiv_id,
            len(captions),
            len(unique_captions),
        )

    return extracted


def filter_images_by_summary(arxiv_id: str, figures: list[dict]) -> int:
    """根据 summary 中的 figures 字段过滤提取的图片/表格。

    用 manifest.json 中的 label 匹配，保留被 AI 总结引用的图片。
    """
    if not figures:
        return 0

    images_dir = paper_dir(arxiv_id) / "images"
    manifest_path = images_dir / "manifest.json"

    if not images_dir.exists() or not manifest_path.exists():
        return 0

    all_files = [f for f in images_dir.iterdir() if f.suffix == ".png"]
    if not all_files:
        return 0

    manifest: dict = json.loads(manifest_path.read_text(encoding="utf-8"))

    # 收集 summary 中引用的所有 Figure/Table ID（归一化）
    referenced_ids: set[str] = set()
    for fig in figures:
        fig_id = fig.get("id", "")
        m = re.match(r"(?:Fig\.?|Figure)\s*(\d+)", fig_id, re.IGNORECASE)
        if m:
            referenced_ids.add(f"Figure {m.group(1)}")
        m2 = re.match(r"Table\s*(\d+)", fig_id, re.IGNORECASE)
        if m2:
            referenced_ids.add(f"Table {m2.group(1)}")

    if not referenced_ids:
        logger.warning("No valid figure/table IDs in summary for %s", arxiv_id)
        return len(all_files)

    # 根据 manifest 的 label 字段匹配
    keep_filenames: set[str] = set()
    for filename, info in manifest.items():
        label = info.get("label", "")
        if label in referenced_ids:
            keep_filenames.add(filename)
            continue
        for ref in info.get("figures", []) + info.get("tables", []):
            if ref in referenced_ids:
                keep_filenames.add(filename)
                break

    if not keep_filenames:
        logger.warning(
            "No manifest matches for %s (refs=%s), keeping all",
            arxiv_id,
            referenced_ids,
        )
        return len(all_files)

    removed = 0
    for f in all_files:
        if f.name not in keep_filenames:
            f.unlink()
            removed += 1

    kept = len(all_files) - removed
    logger.info(
        "Filtered images for %s: kept %d, removed %d (refs=%s)",
        arxiv_id,
        kept,
        removed,
        referenced_ids,
    )
    return kept