"""PDF 图片与表格提取。

DocLayout-YOLO 检测 figure/table 内容区域 → PDF 文本流定位 caption → 只渲染配到
Figure/Table 标题的，用 caption 自带权威 ID 命名。没配到标题的（Algorithm 伪代码、
无编号附录表、DocLayout 误检碎片）一律过滤，不输出。

caption 定位用 PDF 文本而非 DocLayout 的 caption box —— 后者检测不稳（多行标题只
框一行→截断、漏检→无标题、配对错误→串台）。page.get_text("dict") 找以
"Figure N"/"Table N" 开头的文本块：文本块天然含完整多行标题，且其 ID 即论文实际
编号，直接命名规避串台。figure 标题优先在下方、table 标题优先在上方配对。
"""

from __future__ import annotations

import json
import logging
import re
from dataclasses import dataclass
from pathlib import Path

import pymupdf

from app.services.layout_detector import LayoutBox, detect_page_layout
from app.services.pdf_downloader import paper_dir
from app.utils import PAPERS_DIR, TMP_DIR

logger = logging.getLogger(__name__)

# 截图区域的外边距（单位: pt）
_REGION_PADDING = 5
# 渲染倍率（3x 保证清晰度）
_RENDER_ZOOM = 3
# 相邻 box 聚类间距（单位: pt）— 同一 figure/table 的碎片间距通常 < 15pt
_CLUSTER_GAP = 15
# 最小 bbox 面积（单位: pt²）— 过滤 icon/logo 等微小误检
_MIN_BOX_AREA = 2000
# caption 文本块与 figure/table 内容块的最大垂直距离（单位: pt）
_CAPTION_MATCH_DISTANCE = 120
# 游离碎片（配不到 caption）并入紧邻已配 cluster 的最大垂直间距（单位: pt）
# 容下多面板图的 "(a)/(b)" 子图标占位（实测 Figure 4 两面板间距 36pt）
_ABSORB_GAP = 60
# caption 开头标记：Figure 3 / Fig. 3 / Table C1 / Figure 3.5 等（大小写均可）
# 编号 = 数字开头 或 字母+数字（附录 C1）；行首匹配，规避正文 "see Table 3" 引用。
# 否定前瞻再排除多图引用型正文 —— "Figure 14 and 15 show..." / "Figure 1 to 3" /
# "Table 2, 3" 这类引用多张图表的句子不是独立标题，真标题从不引用多个编号。
_CAPTION_HEAD_RE = re.compile(
    r"^\s*(Figure|Fig\.?|Table)\b\.?\s+([0-9][0-9A-Za-z.]*|[A-Z]\d[0-9A-Za-z.]*)"
    r"(?![0-9A-Za-z.])"  # 编号须是完整 token，防止 "Figure 14" 回溯成 "Figure 1" 逃逸
    r"(?!\s*(?:and|to|through|vs\.?|&)\s+\d)"  # "Figure 3 and 4" / "Figure 1 to 3"
    r"(?!\s*,\s*\d)",  # "Figure 3, 4"
    re.IGNORECASE,
)


# ── Box 聚类 ─────────────────────────────────────────────────────────


class _BoxCluster:
    """合并后的布局区域（由一个或多个相邻 LayoutBox 组成）。"""

    __slots__ = ("x0", "y0", "x1", "y1", "boxclass")

    def __init__(self, boxes: list):
        self.x0 = min(b.x0 for b in boxes)
        self.y0 = min(b.y0 for b in boxes)
        self.x1 = max(b.x1 for b in boxes)
        self.y1 = max(b.y1 for b in boxes)
        self.boxclass = boxes[0].boxclass


@dataclass(frozen=True)
class _CaptionBlock:
    """从 PDF 文本流提取的标题块：自带权威 ID、完整多行文本、精确 bbox。"""

    id: str  # "Figure 3" / "Table C1"
    kind: str  # "figure" | "table"
    text: str  # 完整多行标题文本
    bbox: list[float]  # [x0, y0, x1, y1]


def _cluster_to_box(cluster: _BoxCluster) -> list[float]:
    return [
        round(float(cluster.x0), 1),
        round(float(cluster.y0), 1),
        round(float(cluster.x1), 1),
        round(float(cluster.y1), 1),
    ]


def _cluster_boxes(boxes: list, gap: float = _CLUSTER_GAP) -> list[_BoxCluster]:
    """将相邻的同类型 box 合并为聚类。"""
    if not boxes:
        return []

    n = len(boxes)
    parent = list(range(n))

    def find(x: int) -> int:
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a: int, b: int) -> None:
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[ra] = rb

    for i in range(n):
        bi = boxes[i]
        for j in range(i + 1, n):
            bj = boxes[j]
            if bi.boxclass != bj.boxclass:
                continue
            h_gap = max(0.0, max(bi.x0, bj.x0) - min(bi.x1, bj.x1))
            v_gap = max(0.0, max(bi.y0, bj.y0) - min(bi.y1, bj.y1))
            h_overlap = bi.x1 > bj.x0 - gap and bj.x1 > bi.x0 - gap
            v_overlap = bi.y1 > bj.y0 - gap and bj.y1 > bi.y0 - gap
            if (h_gap <= gap and v_overlap) or (v_gap <= gap and h_overlap):
                union(i, j)

    groups: dict[int, list] = {}
    for i in range(n):
        groups.setdefault(find(i), []).append(boxes[i])

    return [_BoxCluster(members) for members in groups.values()]


def _find_caption_blocks(page) -> list[_CaptionBlock]:
    """从页面文本流提取以 "Figure N"/"Table N"/"Fig. N" 开头的标题块。

    用 PDF 文本而非 DocLayout caption box：文本块天然含完整多行标题，
    且其 ID 即论文实际编号（如 "Table C1"），权威且不依赖模型检测。
    """
    try:
        d = page.get_text("dict")
    except Exception:
        return []

    results: list[_CaptionBlock] = []
    for block in d.get("blocks", []):
        if block.get("type") != 0:  # 仅文本块
            continue
        lines = block.get("lines", [])
        if not lines:
            continue
        line_texts = [
            "".join(span.get("text", "") for span in line.get("spans", []))
            for line in lines
        ]
        joined = " ".join(t.strip() for t in line_texts if t.strip())
        m = _CAPTION_HEAD_RE.match(joined)
        if not m:
            continue
        kind_word, num = m.group(1), m.group(2)
        is_table = kind_word.lower().startswith("table")
        bbox = block.get("bbox")
        if not bbox or len(bbox) != 4:
            continue
        full_text = joined
        results.append(
            _CaptionBlock(
                id=f"{'Table' if is_table else 'Figure'} {num}",
                kind="table" if is_table else "figure",
                text=full_text,
                bbox=[float(v) for v in bbox],
            )
        )
    return results


def _pair_caption_blocks(
    content_clusters: list[_BoxCluster],
    caption_blocks: list[_CaptionBlock],
) -> dict[int, list[int]]:
    """每个 caption 配对其垂直 span 内的所有同类型 cluster（支持复合图/子表）。

    不预设标题在内容上方还是下方 —— figure 惯例标题在下、table 惯例标题在上，
    但不少论文反向排版（table 标题在表下方）；用方向作硬约束或加错向惩罚会把
    "同页相邻两张表" 错并（标题居中的那张吞掉邻居）。改为上下两侧平等地按垂直
    距离打分，每个 cluster 唯一归属最近的同类 caption，但一个 caption 可被多个
    cluster 共享 —— 这样一张被 DocLayout 切成多个稀疏子框的复合图/复合表，能
    整体配到它的主标题（而非只截其中一个子图/子表）。

    Returns:
        caption_idx → [cluster_idx, ...]，按 cluster 在页面上的位置排序，
        保证合并/渲染顺序稳定。
    """
    candidates: list[tuple[float, int, int]] = []
    for c_idx, content in enumerate(content_clusters):
        cluster_kind = "figure" if content.boxclass == "picture" else "table"
        for b_idx, cap in enumerate(caption_blocks):
            if cap.kind != cluster_kind:  # 类型过滤：防 figure↔table 串台
                continue
            cx0, cy0, cx1, cy1 = cap.bbox
            h_overlap = min(content.x1, cx1) - max(content.x0, cx0)
            min_width = min(content.x1 - content.x0, cx1 - cx0)
            if min_width <= 0 or h_overlap < min_width * 0.25:
                continue
            if cy1 <= content.y0:  # 标题在内容上方
                v_gap = content.y0 - cy1
            elif cy0 >= content.y1:  # 标题在内容下方
                v_gap = cy0 - content.y1
            else:
                continue  # 重叠，跳过
            if v_gap > _CAPTION_MATCH_DISTANCE:
                continue
            candidates.append((v_gap, c_idx, b_idx))

    # cluster 唯一归属最近的 caption；caption 可被多个 cluster 共享（复合图/子图）
    cluster_to_caption: dict[int, int] = {}
    for _score, c_idx, b_idx in sorted(candidates):
        if c_idx in cluster_to_caption:
            continue
        cluster_to_caption[c_idx] = b_idx

    # 聚合 caption → clusters，按页面位置排序保证稳定的合并/渲染顺序
    caption_to_clusters: dict[int, list[int]] = {}
    for c_idx, b_idx in cluster_to_caption.items():
        caption_to_clusters.setdefault(b_idx, []).append(c_idx)
    for b_idx in caption_to_clusters:
        caption_to_clusters[b_idx].sort(
            key=lambda i: (content_clusters[i].y0, content_clusters[i].x0)
        )
    return caption_to_clusters


def _absorb_stragglers(
    clusters: list[_BoxCluster],
    caption_blocks: list[_CaptionBlock],
    caption_matches: dict[int, list[int]],
) -> None:
    """把配不到 caption、却紧邻某已配 cluster 的同类型游离碎片并入该 cluster 所属 caption。

    多面板图常把主 caption 只放在最下方，离 caption 过远（> _CAPTION_MATCH_DISTANCE）
    的上方面板会被正常配对漏掉（如 Figure 4 的 (a) 子图距主标题 270pt）。这里把它们
    并入紧邻的、已归属某 caption 的同类型 cluster，两道护栏避免误并：

      1. 垂直紧邻（≤ _ABSORB_GAP）且水平重叠（同一列）——排除不同列的无关图；
      2. 两者之间不得夹其他 caption ——有 caption 即另一张图/表的边界（如 Table 7/8
         之间夹 Table 7 标题），不并，从而不破坏表格的独立配对。
    """
    # cluster_idx → caption_idx（已配 cluster 的反向索引）
    cluster_to_cap: dict[int, int] = {}
    for cap_idx, idxs in caption_matches.items():
        for c_idx in idxs:
            cluster_to_cap[c_idx] = cap_idx

    # 多趟扫描直到稳定：游离面板可能链式排列（上图→中图→下图→caption），
    # 单趟只能吸收紧邻已配 cluster 的那一层；中图被并入后才轮到上图，故需重复。
    changed = True
    while changed:
        changed = False
        for u_idx, u in enumerate(clusters):
            if u_idx in cluster_to_cap:
                continue  # 已配，无需吸收
            best_cap: int | None = None
            best_gap: float | None = None
            for p_idx, cap_idx in cluster_to_cap.items():
                p = clusters[p_idx]
                if p.boxclass != u.boxclass:
                    continue  # 类型不同（figure vs table）不并
                gap = max(0.0, max(u.y0, p.y0) - min(u.y1, p.y1))
                if gap > _ABSORB_GAP:
                    continue  # 护栏 1：垂直不紧邻
                if min(u.x1, p.x1) - max(u.x0, p.x0) <= 0:
                    continue  # 护栏 1：水平不重叠，非同一列
                # 护栏 2：u 与 p 之间的垂直区间不得夹任何 caption
                between_lo = min(u.y1, p.y1)
                between_hi = max(u.y0, p.y0)
                if any(
                    not (cb.bbox[3] < between_lo or cb.bbox[1] > between_hi)
                    for cb in caption_blocks
                ):
                    continue
                if best_gap is None or gap < best_gap:
                    best_gap, best_cap = gap, cap_idx
            if best_cap is not None:
                caption_matches[best_cap].append(u_idx)
                cluster_to_cap[u_idx] = best_cap  # 标记已并入，供后续趟链式吸收
                changed = True


# ── Phase 1: 检测 + 渲染 ──────────────────────────────────────────────


def _render_box(
    page,
    box: _BoxCluster,
    images_dest: Path,
    filename: str,
    cap_type: str,
    page_num: int,
    caption_bbox: list[float] | None = None,
) -> bool:
    """渲染单个 box 区域并保存 JPEG，成功返回 True。

    若提供 caption_bbox，则将内容与标题区域合并后一起截取，
    使同一张截图同时包含图/表及其完整标题。
    """
    page_width = page.rect.width
    page_height = page.rect.height
    x0, y0, x1, y1 = box.x0, box.y0, box.x1, box.y1
    if caption_bbox is not None:
        cx0, cy0, cx1, cy1 = caption_bbox
        x0 = min(x0, cx0)
        y0 = min(y0, cy0)
        x1 = max(x1, cx1)
        y1 = max(y1, cy1)
    clip = pymupdf.Rect(
        max(0, x0 - _REGION_PADDING),
        max(0, y0 - _REGION_PADDING),
        min(page_width, x1 + _REGION_PADDING),
        min(page_height, y1 + _REGION_PADDING),
    )
    mat = pymupdf.Matrix(_RENDER_ZOOM, _RENDER_ZOOM)
    try:
        pix = page.get_pixmap(matrix=mat, clip=clip)
    except Exception:
        return False

    (images_dest / filename).write_bytes(pix.tobytes("jpeg", jpg_quality=92))
    return True


def _process_page(
    doc,
    page_idx: int,
    page_boxes: list[LayoutBox],
    images_dest: Path,
    manifest: dict,
    seen_labels: set,
    arxiv_id: str,
) -> int:
    """处理单页：检测内容 box → 文本定位 caption → 只渲染配到标题的。

    配到 Figure/Table caption 的 box 用 caption 自带 ID 命名（figure_3.jpg）；
    同一 caption 的多个 cluster（复合图/子图被 DocLayout 切散）合并 bbox 整张截取；
    没配到标题的（Algorithm 伪代码、无编号附录表、误检碎片）一律过滤，不输出。
    """
    page = doc[page_idx]
    page_num = page_idx + 1

    # 收集本页 figure/table 内容 box（跳过极小区域；caption 改由文本定位，不收 box）
    raw_boxes = []
    for box in page_boxes:
        if box.boxclass in ("table", "picture"):
            w = box.x1 - box.x0
            h = box.y1 - box.y0
            if w < 20 or h < 20 or w * h < _MIN_BOX_AREA:
                continue
            raw_boxes.append(box)

    if not raw_boxes:
        return 0

    # 聚类：将同一 figure/table 的碎片 box 合并；用 PDF 文本定位 caption
    clusters = _cluster_boxes(raw_boxes)
    caption_blocks = _find_caption_blocks(page)
    # caption_idx → [cluster_idx, ...]；一个 caption 可含多个子图 cluster（复合图）
    caption_matches = _pair_caption_blocks(clusters, caption_blocks)
    # 把距主 caption 过远的游离面板（如多面板图上方的 (a) 子图）并入紧邻的已配 cluster
    _absorb_stragglers(clusters, caption_blocks, caption_matches)

    extracted = 0
    for cap_idx, cluster_indices in caption_matches.items():
        cap = caption_blocks[cap_idx]
        if cap.id in seen_labels:
            continue  # 同一图表被 DocLayout 切成多块重复检测，跳过后续
        seen_labels.add(cap.id)

        # 同一 caption 的所有 cluster 合并 bbox，复合图/子图整张截取
        members = [clusters[i] for i in cluster_indices]
        merged = _BoxCluster(members)

        filename = f"{cap.id.replace(' ', '_').lower()}.jpg"
        if not _render_box(
            page,
            merged,
            images_dest,
            filename,
            cap.kind,
            page_num,
            caption_bbox=cap.bbox,
        ):
            continue

        info = {
            "page": page_num,
            "type": cap.kind,
            "label": cap.id,
            "box": _cluster_to_box(merged),
            "caption_text": cap.text[:500],
            "caption_box": cap.bbox,
            "caption_source": "text",
        }
        if len(members) > 1:
            info["subfigure_count"] = len(members)
        manifest[filename] = info
        extracted += 1

    return extracted


# ── Phase 1 核心入口 ───────────────────────────────────────────────────


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

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

    Returns:
        提取的图片数量
    """
    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.iterdir():
        if old_file.suffix.lower() in (".png", ".jpg", ".jpeg"):
            old_file.unlink()
    if (images_dest / "manifest.json").exists():
        (images_dest / "manifest.json").unlink()

    with pymupdf.open(str(pdf_path)) as doc:
        extracted = 0
        manifest: dict[str, dict] = {}
        seen_labels: set[str] = set()

        for page_idx in range(doc.page_count):
            try:
                page_boxes = detect_page_layout(doc[page_idx])
                extracted += _process_page(
                    doc,
                    page_idx,
                    page_boxes,
                    images_dest=images_dest,
                    manifest=manifest,
                    seen_labels=seen_labels,
                    arxiv_id=arxiv_id,
                )
            except Exception:
                logger.warning(
                    "Failed to process page %d for %s",
                    page_idx + 1,
                    arxiv_id,
                    exc_info=True,
                )
                continue

    # 保存 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",
            extracted,
            arxiv_id,
        )

    return extracted


# ── Figure ↔ Image 关联 ────────────────────────────────────────────────


def _normalize_figure_id(raw_id: str) -> str:
    """归一化 Figure/Table ID：'Figure 1'/'Fig.1' → 'Figure 1'。"""
    m = re.match(r"(?:Fig\.?|Figure)\s*(\d+)", raw_id, re.IGNORECASE)
    if m:
        return f"Figure {m.group(1)}"
    m2 = re.match(r"Table\s*(\d+)", raw_id, re.IGNORECASE)
    if m2:
        return f"Table {m2.group(1)}"
    return raw_id


def _is_figure_type(fig_id: str) -> bool:
    """判断是否为 Figure 类型（非 Table）。"""
    return not re.match(r"Table\s*(\d+)", fig_id, re.IGNORECASE)


def _image_sort_key(name: str) -> tuple[int, int]:
    """按文件名中的编号排序提取的图片。"""
    # 新格式：figure_1.jpg, table_1.jpg
    m = re.search(r"(?:figure|table)_(\d+)", name)
    if m:
        return (0, int(m.group(1)))
    return (0, 0)


def link_figures_with_images(
    figures: list[dict], images: list[dict], arxiv_id: str
) -> list[dict]:
    """将 summary figures 元数据与提取的图片文件关联。

    策略：
    1. 优先用 manifest.json 的 label 做 ID 精确匹配
    2. 未匹配的 figure 用序号兜底：第 N 个 Figure → 第 N 张提取图
    """
    if not figures or not images:
        return figures

    manifest_path = PAPERS_DIR / arxiv_id / "images" / "manifest.json"

    # ── 策略 1：manifest ID 精确匹配 ──
    id_to_url: dict[str, str] = {}
    if manifest_path.exists():
        try:
            manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
        except (ValueError, TypeError):
            manifest = {}
        for filename, info in manifest.items():
            url = f"/papers/{arxiv_id}/images/{filename}"
            # 优先用 label 字段（新格式）
            label = info.get("label", "")
            if label:
                id_to_url[label] = url
            # 也兼容 figures/tables 列表（旧格式）
            for fig_id in info.get("figures", []) + info.get("tables", []):
                if fig_id not in id_to_url:
                    id_to_url[fig_id] = url

    for fig in figures:
        raw_id = fig.get("id", "")
        normalized = _normalize_figure_id(raw_id)
        if normalized in id_to_url:
            fig["image_url"] = id_to_url[normalized]

    # ── 策略 2：序号兜底（manifest 匹配不到时） ──
    unmatched = [f for f in figures if not f.get("image_url")]
    if not unmatched:
        return figures

    # 已被策略 1（精确匹配）占用的图片不参与兜底，否则会把已正确归属的图复用给
    # 别的条目（如缺失的 Table 4 误链到 Table 1 的截图）。
    assigned_urls = {f["image_url"] for f in figures if f.get("image_url")}

    # 按类型分流：Figure vs Table
    fig_type_unmatched = [f for f in unmatched if _is_figure_type(f.get("id", ""))]
    table_type_unmatched = [
        f for f in unmatched if not _is_figure_type(f.get("id", ""))
    ]

    # 剩余未占用图片按类型分流，按文件名中的编号排序
    fig_images = sorted(
        [
            img
            for img in images
            if "table" not in img["name"].lower() and img["url"] not in assigned_urls
        ],
        key=lambda img: _image_sort_key(img["name"]),
    )
    table_images = sorted(
        [
            img
            for img in images
            if "table" in img["name"].lower() and img["url"] not in assigned_urls
        ],
        key=lambda img: _image_sort_key(img["name"]),
    )

    for i, fig in enumerate(fig_type_unmatched):
        if i < len(fig_images):
            fig["image_url"] = fig_images[i]["url"]

    for i, fig in enumerate(table_type_unmatched):
        if i < len(table_images):
            fig["image_url"] = table_images[i]["url"]

    return figures