#!/usr/bin/env python3
"""DrissionPage批量抓取成都区县组织部/党校2026年预算 - v3"""
import json, time, os, re, sys
from DrissionPage import ChromiumPage, ChromiumOptions

OUT = "/data/www/files/budget_pdfs/chengdu_districts"
SS = "/data/www/files/budget_pdfs/screenshots"
RES = "/data/www/files/budget_data/cd_districts_v3.json"
os.makedirs(OUT, exist_ok=True)
os.makedirs(SS, exist_ok=True)

DISTS = [
    ("青羊区", "www.qingyang.gov.cn", "qingyang"),
    ("武侯区", "www.wuhou.gov.cn", "wuhou"),
    ("成华区", "www.chenghua.gov.cn", "chenghua"),
    ("新都区", "www.xindu.gov.cn", "xindu"),
    ("温江区", "www.wenjiang.gov.cn", "wenjiang"),
    ("郫都区", "www.pidu.gov.cn", "pidu"),
    ("新津区", "www.xinjin.gov.cn", "xinjin"),
    ("青白江区", "www.qbj.gov.cn", "qingbaijiang"),
    ("都江堰市", "www.djy.gov.cn", "dujiangyan"),
    ("彭州市", "www.pengzhou.gov.cn", "pengzhou"),
    ("邛崃市", "www.qionglai.gov.cn", "qionglai"),
    ("崇州市", "www.chongzhou.gov.cn", "chongzhou"),
    ("简阳市", "www.jianyang.gov.cn", "jianyang"),
    ("金堂县", "www.jintang.gov.cn", "jintang"),
    ("大邑县", "www.dayi.gov.cn", "dayi"),
    ("蒲江县", "www.pujiang.gov.cn", "pujiang"),
]

def abs_url(href, base, code):
    if href.startswith("http"):
        return href
    if href.startswith("/"):
        return f"https://{base}{href}"
    return f"https://{base}/{code}/{href}"

results = {}

co = ChromiumOptions()
co.set_argument("--no-sandbox")
co.set_argument("--disable-dev-shm-usage")
co.set_argument("--disable-gpu")
co.set_argument("--window-size=1920,1080")
co.headless(True)
page = ChromiumPage(co)
print("浏览器启动", flush=True)

for name, domain, code in DISTS:
    print(f"\n=== {name} ===", flush=True)
    base = f"https://{domain}"
    dr = {"组织部": [], "党校": [], "screenshots": [], "status": "scanning"}

    # 1) 首页建立WAF
    try:
        page.get(base + "/")
        time.sleep(4)
        t = (page.title or "")[:50]
        print(f"  首页OK: {t}", flush=True)
    except Exception as e:
        print(f"  首页失败: {e}", flush=True)
        dr["status"] = "unreachable"
        results[name] = dr
        continue

    # 2) 尝试多个栏目路径
    cols = ["c110728","c110729","c100043","c111709","c115761","c106855","c106854"]
    for col in cols:
        for pn in range(1, 5):
            suf = f"_{pn}" if pn > 1 else ""
            url = f"{base}/{code}/{col}/list{suf}.shtml"
            try:
                page.get(url)
                time.sleep(1.8)
                html = page.html or ""
                if len(html) < 100:
                    continue
                links = re.findall(r'<a[^>]*href=["\']([^"\']+)["\'][^>]*>([^<]{4,100})</a>', html)
                for href, txt in links:
                    txt = txt.strip()
                    if not txt:
                        continue
                    for kw in ["组织部", "党校"]:
                        if kw in txt and ("2026" in txt or "预算" in txt or "2026" in href):
                            fu = abs_url(href, domain, code)
                            ex = [x["title"] for x in dr[kw]]
                            if txt not in ex:
                                dr[kw].append({"title": txt, "url": fu})
                                print(f"  ✅ {kw}: {txt[:55]}", flush=True)
            except:
                continue
    
    # 3) 如果还没找到，用搜索
    if not dr["组织部"] and not dr["党校"]:
        for kw in ["组织部", "党校"]:
            try:
                surl = f"{base}/search.html?q={kw}+2026+预算"
                page.get(surl)
                time.sleep(2)
                html = page.html or ""
                links = re.findall(r'<a[^>]*href=["\']([^"\']+)["\'][^>]*>([^<]{4,100})</a>', html)
                for href, txt in links:
                    txt = txt.strip()
                    if kw in txt and ("2026" in txt or "预算" in txt):
                        fu = abs_url(href, domain, code)
                        ex = [x["title"] for x in dr[kw]]
                        if txt not in ex:
                            dr[kw].append({"title": txt, "url": fu})
                            print(f"  🔍 {kw}: {txt[:55]}", flush=True)
            except:
                continue

    # 4) 截图 + 下载附件
    for kw in ["组织部", "党校"]:
        for item in dr[kw][:2]:
            try:
                page.get(item["url"])
                time.sleep(2)
                sn = f"{code}_{kw}.png"
                page.get_screenshot(path=f"{SS}/{sn}", full_page=True)
                dr["screenshots"].append({"file": f"budget_pdfs/screenshots/{sn}", "title": item["title"]})
                print(f"  📸 {sn}", flush=True)

                html = page.html or ""
                atts = re.findall(r'href=["\']([^"\']*\.(?:pdf|xlsx?|docx?))["\']', html, re.I)
                dl = []
                for au in atts[:3]:
                    fau = abs_url(au, domain, code)
                    fn = au.split("/")[-1].split("?")[0]
                    sp = f"{OUT}/{code}_{kw}_{fn}"
                    try:
                        js = f"return await fetch('{fau}').then(r=>r.arrayBuffer()).then(b=>Array.from(new Uint8Array(b)))"
                        data = page.run_js(js)
                        if data and len(data) > 300:
                            with open(sp, "wb") as f:
                                f.write(bytes(data))
                            sz = os.path.getsize(sp)
                            dl.append({"url": fau, "local": f"budget_pdfs/chengdu_districts/{code}_{kw}_{fn}", "size": sz})
                            print(f"  📎 {fn} ({sz//1024}KB)", flush=True)
                    except:
                        pass
                item["attachments"] = dl
            except Exception as e:
                print(f"  详情失败: {e}", flush=True)

    dr["status"] = "completed"
    results[name] = dr
    with open(RES, "w") as f:
        json.dump(results, f, ensure_ascii=False, indent=2)
    print(f"  💾 已保存", flush=True)

try:
    page.quit()
except:
    pass

org = sum(len(r.get("组织部",[])) for r in results.values())
dx = sum(len(r.get("党校",[])) for r in results.values())
fl = sum(len(a) for r in results.values() for k in ["组织部","党校"] for i in r.get(k,[]) for a in i.get("attachments",[]))
ss = sum(len(r.get("screenshots",[])) for r in results.values())
print(f"\n===== 完成 =====", flush=True)
print(f"组织部: {org}条, 党校: {dx}条, 文件: {fl}个, 截图: {ss}张", flush=True)