#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""PDF → 内容模型 v2：按块合并段落、过滤页眉页脚、导出图片 PNG"""
import json
import os
import re
import sys

import fitz


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


HEADER_PAT = [
    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{4,}_\d+.*\.PDF$'),
    re.compile(r'^0?_EN-\d+\.PDF$'),
]


def is_noise(line):
    s = clean(line)
    if not s:
        return True
    return any(p.match(s) for p in HEADER_PAT)


def join_lines(lines):
    out = ''
    for ln in lines:
        ln = clean(ln)
        if not ln:
            continue
        if not out:
            out = ln
            continue
        # 中文直接拼，英文/数字加空格
        if re.search(r'[\u4e00-\u9fff]$', out) or re.match(r'^[\u4e00-\u9fff]', ln):
            out += ln
        else:
            out += ' ' + ln
    return out


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


def extract(pdf_path, outdir, prefix, max_pages=None):
    os.makedirs(outdir, exist_ok=True)
    doc = fitz.open(pdf_path)
    pages = []
    for pno in range(len(doc)):
        if max_pages and pno >= max_pages:
            break
        page = doc[pno]
        raw = page.get_text('dict')
        sizes, blocks = [], []
        for b in raw['blocks']:
            if b.get('type') != 0:
                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 = join_lines(lines)
            if is_noise(txt):
                continue
            blocks.append({'bbox': list(b['bbox']), 'text': txt, 'size': round(sz, 1)})
            sizes.append(sz)

        body = sorted(sizes)[len(sizes) // 2] if sizes else 11

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

        # 元素（文本块 + 图片）；表格单独由 pdfplumber 处理（此版先以文本块呈现）
        items = []
        for b in blocks:
            items.append({'type': 'h' if heading_like(b['text'], b['size'], body) else 'p',
                          'text': b['text'], 'size': b['size'], 'bbox': b['bbox']})
        for im in imgs:
            items.append(im)
        # 表格区域文本已包含在文本块中，表格结构由第二遍 pdfplumber 覆盖
        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]:
        for it in p['items'][:10]:
            print(' ', it['type'], it.get('size', ''), (it.get('text') or it.get('path', ''))[:70])