# -*- coding: utf-8 -*-
"""把 WPS 转出的 PDF 里卷商广告页眉/页脚涂掉（PyMuPDF 真删文字，不是盖白块）。
   用法：~/.venvs/pdf-tools/bin/python strip_ads.py [--dry] [--limit N] [--only 路径子串]
   规矩：
   - 只动我们自己转出的 PDF（块目录里的 + 归位到源旁的同一份），卷商原本自带的 PDF 不碰。
   - 只删含卷商关键词（KEY：微信号/店名，不是泛泛的「微信」）的**整个文本块**，且只在页眉区（上 12%）或页脚区（下 12%）——
     题面里的「微信收付款」「微信钱包账单」这类正文绝不会被碰（关键词不命中，位置也不在带内）。
   - 每份处理后：页数不变、pdftotext 再搜短语=0，才覆盖；否则保留原样并记入 failed。"""
import json, os, re, sys, subprocess, shutil, hashlib
import fitz  # PyMuPDF

HERE = os.path.dirname(os.path.abspath(__file__))
ROOT = os.path.dirname(HERE)
PHRASES = [
    '微信：cs5311742016 主营各地名校期中期末真题卷及小初高全套教材教案讲义知识点总结',
    '微信：cs5311742016 主营全国各地小初高名校月考期中期末真题卷',
    '微信：cs5311742016',
    '【淘宝搜索店铺：中小学教辅资源店 微信：mlxt2022】',
    '【中小学教辅资源店 微信：mlxt2022】',
    '淘宝：中小学教辅资源店 微信：mlxt2022',
    '【淘宝：中小学教辅资源店',
    '淘宝：中小学教辅资源店',
    '微信：mlxt2022】', '微信：mlxt2022',
    '更多资料添加微信号：DEM2008 淘宝搜索店铺：优尖升教育 网址：shop492842749.taobao.com',
    '更多资料添加微信号：DEM2008',
    '淘宝搜索店铺：优尖升教育 网址：shop492842749.taobao.com',
    '淘宝搜索店铺：优尖升教育', 'shop492842749.taobao.com',
    '认准淘宝店：真学子资源店', '微信：2496342225',
]
KEY = re.compile(r'cs5311742016|mlxt2022|DEM2008|优尖升|真学子|2496342225|中小学教辅资源店|shop492842749')
BAND = 0.12   # 页眉/页脚区占页高比例

dry = '--dry' in sys.argv
limit = int(sys.argv[sys.argv.index('--limit') + 1]) if '--limit' in sys.argv else 0
only = sys.argv[sys.argv.index('--only') + 1] if '--only' in sys.argv else ''

def _md5(f):
    return hashlib.md5(open(f, 'rb').read()).hexdigest()

def targets():
    """(块目录 PDF, 源旁 PDF 或 None)；源旁那份只有在与块目录那份 **md5 完全相同**（即我们放的）时才算。
       ⚠ 不能用大小判：卷商 PDF 多半也是 WPS 导出同一份 docx，大小相同、只差元数据（1,357 份撞大小）。"""
    out = []
    for b in ['_pdf_batch', '_pdf_batch2']:
        for m in json.load(open(os.path.join(ROOT, b, 'manifest.json'), encoding='utf-8')):
            p = os.path.join(ROOT, b, m['chunk'], os.path.splitext(m['name'])[0] + '.pdf')
            d = os.path.splitext(m['src'])[0] + '.pdf'
            same = os.path.exists(d) and os.path.getsize(d) == os.path.getsize(p) and _md5(d) == _md5(p)
            out.append((p, d if same else None))
    return out

def text_has_ad(p):
    o = subprocess.run(['pdftotext', p, '-'], capture_output=True, text=True, timeout=60).stdout
    return bool(KEY.search(o))

def redact(src, dst):
    """按**文本块**涂：页眉/页脚区里含卷商关键词的整块一起删（短语表只是备忘，卷商页眉常无空格、拆行，按短语搜会留尾巴）。
       正文区含关键词的块只计数不动。"""
    doc = fitz.open(src); n_pages = len(doc); n_rect = 0; skipped_body = 0
    for page in doc:
        h = page.rect.height; rects = []
        for b in page.get_text('blocks'):
            if b[6] != 0 or not KEY.search(b[4]): continue      # 只看文字块
            r = fitz.Rect(b[:4])
            if r.y1 <= h * BAND or r.y0 >= h * (1 - BAND): rects.append(r)
            else: skipped_body += 1
        for r in rects:
            page.add_redact_annot(r, fill=(1, 1, 1)); n_rect += 1
        if rects:
            page.apply_redactions(images=fitz.PDF_REDACT_IMAGE_NONE)
    doc.save(dst, garbage=3, deflate=True); doc.close()
    return n_pages, n_rect, skipped_body

hits = {h[0] for h in json.load(open(os.path.join(HERE, 'ad_scan.json'), encoding='utf-8'))}
done = failed = 0; body_skips = 0; log = []
tmp = os.path.join(HERE, '_strip_tmp.pdf')
for p, d in targets():
    rel = os.path.relpath(p, ROOT)
    if rel not in hits: continue
    if only and only not in p and (not d or only not in d): continue
    if limit and done + failed >= limit: break
    try:
        if not text_has_ad(p):
            log.append({'pdf': rel, 'src_pdf': d, 'ok': True, 'skip': 'no vendor key'}); continue
        n0 = len(fitz.open(p))
        n_pages, n_rect, sb = redact(p, tmp); body_skips += sb
        ok = (n_pages == n0) and n_rect > 0 and not text_has_ad(tmp)
        if not ok:
            failed += 1; log.append({'pdf': rel, 'src_pdf': d, 'ok': False, 'rects': n_rect, 'pages': (n0, n_pages), 'still_ad': text_has_ad(tmp)})
            continue
        if not dry:
            shutil.copyfile(tmp, p)
            if d: shutil.copyfile(tmp, d)
        done += 1; log.append({'pdf': rel, 'src_pdf': d, 'ok': True, 'rects': n_rect, 'pages': n_pages})
    except Exception as e:
        failed += 1; log.append({'pdf': rel, 'src_pdf': d, 'ok': False, 'err': str(e)[:200]})
if os.path.exists(tmp): os.remove(tmp)
print('%s处理 %d | 失败 %d | 正文区命中被保护 %d' % ('[dry] ' if dry else '', done, failed, body_skips))
for x in [l for l in log if not l['ok']][:10]: print('  ✗', x)
json.dump(log, open(os.path.join(HERE, 'strip_ads_report%s.json' % ('_dry' if dry else '')), 'w', encoding='utf-8'), ensure_ascii=False, indent=1)