"""批量提取cceup PDF中的姓名和电话"""
import sys, os, re, csv, json
sys.path.insert(0, os.path.expanduser("~/gov_crawler/venv/lib/python3.9/site-packages"))

import pymupdf
from rapidocr_onnxruntime import RapidOCR

PDF_DIR = os.path.expanduser("~/Nutstore Files/Downloads_nuc")

# ── 提取规则 ──
PHONE_PAT = re.compile(r'(?<!\d)1[3-9]\d{9}(?!\d)')                # 手机号
LANDLINE_PAT = re.compile(r'0[1-9]\d{1,2}-?\d{7,8}')               # 座机
NAME_PAT = re.compile(
    r'(?:联系人|联\s*系\s*人|项目负责人|法人|法人代表|法定代表人|代表|业主|经办人|'
    r'负责人|项目经理|采购人|招标人|姓名|姓\s*名)[：:\s]*([\u4e00-\u9fff·]{2,6})'
)
WATERMARK_PAT = re.compile(r'\(?\d{10,11}\)?\s*\d{4}[-\~]\d{2}[-\~]\d{2}[\s\d:]*')

# 姓名黑名单（字段名误抓）
NAME_BLACKLIST = {
    '手机','代表','单位注册地址','为外资','姓名及号码','姓名',
    '且分管工程','单位办公地址','注册地址','办公地址','单位名称',
    '联系电话','联系电话、','联系电话手机','联系人手机','联系人电话',
}


def is_valid_name(n):
    n = n.strip()
    if not n or n in NAME_BLACKLIST:
        return False
    if len(n) < 2 or len(n) > 6:
        return False
    if not re.match(r'^[\u4e00-\u9fff·]+$', n):
        return False
    return True


def extract_text(pdf_path):
    doc = pymupdf.open(pdf_path)
    text = '\n'.join(page.get_text() for page in doc)
    doc.close()
    return text


def extract_ocr(pdf_path):
    """RapidOCR备用方案"""
    doc = pymupdf.open(pdf_path)
    ocr = RapidOCR()
    all_lines = []
    for i, page in enumerate(doc):
        pix = page.get_pixmap(dpi=200)
        temp_path = f"/tmp/ocr_p_{os.getpid()}_{i}.png"
        with open(temp_path, "wb") as f:
            f.write(pix.tobytes("png"))
        result, _ = ocr(temp_path)
        os.remove(temp_path)
        if result:
            for _, text, score in result:
                if score > 0.8 and not WATERMARK_PAT.match(text.strip()):
                    all_lines.append(text.strip())
    doc.close()
    return '\n'.join(all_lines)


def process_one(pdf_path, use_ocr=False):
    fname = os.path.basename(pdf_path)
    company = fname.replace('.pdf', '')

    # 阶段1: 文字提取
    text = extract_text(pdf_path)
    phones = list(set(PHONE_PAT.findall(text)))
    landlines = list(set(LANDLINE_PAT.findall(text)))
    raw_names = list(set(NAME_PAT.findall(text)))
    names = [n for n in raw_names if is_valid_name(n)]
    method = 'text'

    # 阶段2: 无电话 → OCR
    if (not phones and not landlines) or use_ocr:
        text2 = extract_ocr(pdf_path)
        phones2 = list(set(PHONE_PAT.findall(text2)))
        landlines2 = list(set(LANDLINE_PAT.findall(text2)))
        raw_names2 = list(set(NAME_PAT.findall(text2)))
        names2 = [n for n in raw_names2 if is_valid_name(n)]
        if phones2 or landlines2:
            phones = phones2 or phones
            landlines = landlines2 or landlines
            names = names2 or names
            method = 'ocr'

    return {
        'company': company,
        'contact_name': ' / '.join(names) if names else '',
        'phone': ' / '.join(phones) if phones else '',
        'landline': ' / '.join(landlines) if landlines else '',
    }


if __name__ == "__main__":
    mode = sys.argv[1] if len(sys.argv) > 1 else 'batch'

    if mode == 'single':
        f = sys.argv[2]
        print(json.dumps(process_one(f), ensure_ascii=False, indent=2))

    elif mode == 'batch':
        pdfs = sorted([f for f in os.listdir(PDF_DIR) if f.endswith('.pdf')])
        results = []
        for i, fname in enumerate(pdfs, 1):
            fpath = os.path.join(PDF_DIR, fname)
            try:
                r = process_one(fpath)
                results.append(r)
                nc = r['contact_name'][:20] if r['contact_name'] else '--'
                np = r['phone'][:20] if r['phone'] else '--'
                print(f"[{i:3d}/{len(pdfs)}] {nc:20s} | {np}")
            except Exception as e:
                print(f"[{i:3d}/{len(pdfs)}] ❌ {fname}: {e}")

        # 汇总
        has_name = sum(1 for r in results if r['contact_name'])
        has_phone = sum(1 for r in results if r['phone'] or r['landline'])
        print(f"\n完成: {len(results)}个PDF, 有姓名{has_name}个, 有电话{has_phone}个")

        # 写JSON
        json_path = os.path.expanduser("~/Desktop/cceup_contacts.json")
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        print(f"JSON: {json_path}")

        # 写CSV
        csv_path = os.path.expanduser("~/Desktop/cceup_contacts.csv")
        with open(csv_path, 'w', newline='', encoding='utf-8-sig') as f:
            w = csv.DictWriter(f, fieldnames=['company', 'contact_name', 'phone', 'landline'])
            w.writeheader()
            w.writerows(results)
        print(f"CSV: {csv_path}")