#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""PDF → 内容模型 v3（表格 + 段落合并 + 图片导出）

- 表格：pdfplumber find_tables（保留行列结构）
- 段落：PyMuPDF blocks，过滤页眉页脚与表格区域，按行距合并
- 图片：页面区域 200dpi 导出（跳过 logo）
"""
import json
import os
import re
import sys

import fitz
import pdfplumber


def clean(t):
    t = (t or '').replace('\u00a0', ' ')
    t = re.sub(r'[ \t]+', ' ', t)
    return t.strip()


NOISE = [
    re.compile(r'^2×FastAmpli Premix-UNG'),
    re.compile(r'^【货号[:：]\s*M2181-4】'),
    re.compile(r'^第\s*\d+\s*页\s*共\s*\d+\s*页'),
    re.compile(r'^CA-10-04-R\d+'),
    re.compile(r'^\d{8}_\d+'),
    re.compile(r'^0?_EN-\d+\.PDF'),
]


def noise(s):
    s = clean(s)
    return (not s) or any(p.match(s) for p in NOISE)


def join(a, b):
    if re.search(r'[\u4e00-\u9fff]$', a) or re.match(r'^[\u4e00-\u9fff]', b):
        return a + b
    return a + ' ' + b


def heading(s, size, body):
    s = clean(s)
    if not s or len(s) > 60:
        return False
    if re.match(r'^\d+(\.\d+)*[\s、\.]', s) and not re.match(r'^\d+\s+\d', s):
        return True
    if re.match(r'^(实验[一二三四五六七八九十]+|附录|附件)', s):
        return True
    if re.match(r'^[一二三四五六七八九十]+[、\.]', s) and len(s) < 30:
        return True
    return size >= body + 1.5 and len(s) <= 40


def extract(pdf_path, imgdir, prefix, max_pages=None):
    os.makedirs(imgdir, exist_ok=True)
    doc = fitz.open(pdf_path)
    pages = []
    with pdfplumber.open(pdf_path) as pl:
        for pno in range(len(doc)):
            if max_pages and pno >= max_pages:
                break
            page, pp = doc[pno], pl.pages[pno]
            W, H = page.rect.width, page.rect.height

            # --- 表格
            tables = []
            for tb in (pp.find_tables() or []):
                data = tb.extract()
                if not data:
                    continue
                rows = [[clean(c) for c in row] for row in data]
                rows = [r for r in rows if any(r)]
                if len(rows) < 1:
                    continue
                tables.append({'type': 'table', 'bbox': list(tb.bbox), 'rows': rows})

            def inside_table(b):
                cy = (b[1] + b[3]) / 2
                for t in tables:
                    x0, y0, x1, y1 = t['bbox']
                    if b[0] >= x0 - 3 and b[2] <= x1 + 3 and y0 - 3 <= cy <= y1 + 3:
                        return True
                return False

            # --- 文本块
            blocks = []
            for b in page.get_text('dict')['blocks']:
                if b.get('type') != 0:
                    continue
                y0, y1 = b['bbox'][1], b['bbox'][3]
                if y1 < 55 or y0 > H - 45:          # 页眉页脚区
                    continue
                lines, sz = [], 0
                for ln in b['lines']:
                    txt = ''.join(sp['text'] for sp in ln['spans'])
                    if txt.strip():
                        lines.append(txt)
                        sz = max(sz, max(sp['size'] for sp in ln['spans']))
                if not lines:
                    continue
                txt = clean(''.join(lines)) if all(re.search(r'[\u4e00-\u9fff]', l) for l in lines) else clean(' '.join(lines))
                if noise(txt) or inside_table(b['bbox']):
                    continue
                blocks.append({'bbox': list(b['bbox']), 'text': txt, 'size': round(sz, 1),
                               'lines': [clean(l) for l in lines if clean(l)]})

            sizes = [b['size'] for b in blocks]
            body = sorted(sizes)[len(sizes) // 2] if sizes else 11

            # 同一块内若每行都是标题（如 4.2 / 4.2.1），拆成独立条目
            split = []
            for b in blocks:
                ls = b.get('lines') or [b['text']]
                if len(ls) > 1 and all(heading(l, b['size'], body) for l in ls):
                    for l in ls:
                        split.append({'bbox': list(b['bbox']), 'text': l, 'size': b['size']})
                else:
                    split.append(b)
            blocks = split

            # 合并相邻正文块（同字号、纵向间距小、上一句未结束）
            blocks.sort(key=lambda x: (x['bbox'][1], x['bbox'][0]))
            merged = []
            for b in blocks:
                if merged:
                    p = merged[-1]
                    gap = b['bbox'][1] - p['bbox'][3]
                    same = abs(b['size'] - p['size']) < 0.6
                    cont = not re.search(r'[。！？；：]\s*$', p['text'])
                    nxt_head = bool(re.match(r'^(\d+(\.\d+)*[\s、\.]|[一二三四五六七八九十]+[、\.]|实验[一二三四五六七八九十]+)', b['text']))
                    nxt_colon = bool(re.match(r'^(?!【)[^：]{2,10}：', b['text']))
                    if same and gap < b["size"] * 1.0 and cont and not nxt_head and not nxt_colon and len(p['text']) < 400:
                        p['text'] = join(p['text'], b['text'])
                        p['bbox'][3] = b['bbox'][3]
                        continue
                merged.append(dict(b))

            items = []
            for t in tables:
                items.append(t)
            for b in merged:
                items.append({'type': 'h' if heading(b['text'], b['size'], body) else 'p',
                              'text': b['text'], 'size': b['size'], 'bbox': b['bbox']})

            # --- 图片
            for im in page.get_images(full=True):
                for r in page.get_image_rects(im[0]):
                    if r.width < 120 or r.height < 60:
                        continue
                    pix = page.get_pixmap(clip=r, dpi=200)
                    fp = os.path.join(imgdir, f'{prefix}_p{pno+1}_{int(r.y0)}_{int(r.x0)}.png')
                    pix.save(fp)
                    items.append({'type': 'img', 'bbox': [r.x0, r.y0, r.x1, r.y1], 'path': fp,
                                  'w': round(r.width), 'h': round(r.height)})

            items.sort(key=lambda x: (x['bbox'][1], x['bbox'][0]))
            pages.append({'page': pno + 1, 'items': items})
    return {'file': os.path.basename(pdf_path), 'pages': pages}


if __name__ == '__main__':
    src, out, prefix = sys.argv[1], sys.argv[2], sys.argv[3]
    mp = int(sys.argv[4]) if len(sys.argv) > 4 else None
    m = extract(src, '/tmp/reformat_imgs', prefix, mp)
    json.dump(m, open(out, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)
    print('pages', len(m['pages']))
    for p in m['pages'][:2]:
        print('-- p', p['page'])
        for it in p['items'][:8]:
            print(' ', it['type'], it.get('size', ''), (it.get('text') or it.get('path', ''))[:70])