#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""PDF → 结构化内容模型（段落/标题/表格/图片）

只做版面元素抽取，不改动任何文字。输出 JSON：
{ "pages": [ { "items": [ {"type":"h|p|table|img", ...} ] } ] }
"""
import json
import os
import re
import sys

import fitz
import pdfplumber


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


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


def extract(pdf_path, max_pages=None):
    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 = doc[pno]
            pwp = pl.pages[pno]

            # 表格：bbox + 单元格
            tables = []
            for tb in (pwp.find_tables() or []):
                data = tb.extract()
                if not data:
                    continue
                rows = [[norm_text(c or '') for c in row] for row in data]
                # 去掉全空行/列
                rows = [r for r in rows if any(c for c in r)]
                if not rows:
                    continue
                tables.append({'type': 'table', 'bbox': list(tb.bbox), 'rows': rows})

            def in_table(b):
                x0, y0, x1, y1 = b
                cy = (y0 + y1) / 2
                for t in tables:
                    tx0, ty0, tx1, ty1 = t['bbox']
                    if tx0 - 2 <= x0 and x1 <= tx1 + 2 and ty0 - 2 <= cy <= ty1 + 2:
                        return True
                return False

            # 文本块（按行）
            blocks = page.get_text('dict')['blocks']
            texts = []
            sizes = []
            for b in blocks:
                if b.get('type') != 0:
                    continue
                lines = []
                sz = 0
                for ln in b['lines']:
                    line = ''.join(sp['text'] for sp in ln['spans'])
                    if line.strip():
                        lines.append(line)
                        sz = max(sz, max(sp['size'] for sp in ln['spans']))
                if not lines:
                    continue
                if in_table(b['bbox']):
                    continue
                texts.append({'type': 'text', 'bbox': list(b['bbox']), 'lines': lines, 'size': sz})
                sizes.extend([sz] * len(lines))

            body_size = sorted(sizes)[len(sizes) // 2] if sizes else 12

            # 图片
            imgs = []
            for im in page.get_images(full=True):
                xref = im[0]
                for r in page.get_image_rects(xref):
                    if r.width < 40 or r.height < 20:
                        continue
                    imgs.append({'type': 'img', 'bbox': [r.x0, r.y0, r.x1, r.y1], 'xref': xref,
                                 'page': pno, 'w': r.width, 'h': r.height})

            items = []
            for t in tables:
                items.append(t)
            for t in texts:
                # 逐行判断标题/正文
                for ln in t['lines']:
                    s = norm_text(ln)
                    if not s:
                        continue
                    items.append({'type': 'h' if is_heading(s, t['size'], body_size) else 'p',
                                  'text': s, 'size': round(t['size'], 1), 'bbox': t['bbox']})
            for im in imgs:
                items.append(im)

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


if __name__ == '__main__':
    src = sys.argv[1]
    out = sys.argv[2]
    mp = int(sys.argv[3]) if len(sys.argv) > 3 else None
    m = extract(src, mp)
    json.dump(m, open(out, 'w', encoding='utf-8'), ensure_ascii=False, indent=1)
    print('pages', len(m['pages']), '→', out)
    for p in m['pages'][:2]:
        print('-- page', p['page'])
        for it in p['items'][:12]:
            if it['type'] == 'table':
                print('  TABLE', len(it['rows']), 'x', len(it['rows'][0]))
            elif it['type'] == 'img':
                print('  IMG', round(it['w']), 'x', round(it['h']))
            else:
                print(' ', it['type'].upper(), it['size'], it['text'][:60])