#!/usr/bin/env python3
"""后台批量维基采集，从 drug_index.json 循环取药"""
import json, time, httpx, os, sys
from pathlib import Path
from bs4 import BeautifulSoup
from urllib.parse import quote

PROJECT = Path(__file__).resolve().parent.parent.parent
os.chdir(PROJECT)
DATA_DIR = PROJECT / "data-pipeline" / "data"

H = {"User-Agent": "PharmacopoeiaAI/1.0 (ResearchBot) Mozilla/5.0", "Accept-Language": "zh-CN,zh;q=0.9"}
SECTION_MAP = {
    "性状": ["性状","外观","物理"],
    "适应症": ["适应","用途","主治","医疗用途","临床应用"],
    "用法与用量": ["用法","用量","剂量","给药","给药方式"],
    "禁忌": ["禁忌","禁用"],
    "不良反应": ["不良反应","副作用","不良事件"],
    "注意事项": ["注意","警告","谨慎"],
    "药理": ["药理","药效","作用机制","药物代谢","药物动力学","药代"],
    "药物相互作用": ["交互作用","药物相互"],
    "贮藏": ["贮藏","储存","保存"],
}

def parse_sections(html):
    soup = BeautifulSoup(html, "lxml")
    body = soup.find("div", class_="mw-parser-output")
    if not body: return {"概述": ""}
    overview_ps = []
    for elem in body.children:
        if elem.name == "div" and "mw-heading" in (elem.get("class") or []): break
        if elem.name == "p":
            t = elem.get_text(" ", strip=True)
            if len(t) > 15: overview_ps.append(t)
    sections = {"概述": " ".join(overview_ps)[:2000]}
    for h_div in body.find_all("div", class_=lambda c: c and "mw-heading" in c):
        span = h_div.find("span", class_="mw-headline")
        heading = (span.get_text() if span else h_div.get_text()).replace("[编辑]","").strip()
        if len(heading) < 2 or len(heading) > 60: continue
        if any(s in heading for s in ["参见","参考","外部","延伸","来源","引用","导航","链接","注释"]): continue
        matched_key = None
        for std_key, patterns in SECTION_MAP.items():
            if any(p in heading for p in patterns): matched_key = std_key; break
        texts = []; nxt = h_div.find_next_sibling()
        while nxt:
            if nxt.name == "div" and "mw-heading" in (nxt.get("class") or []): break
            if nxt.name in ("p","ul","ol"):
                t = nxt.get_text(" ", strip=True)
                if len(t) > 10: texts.append(t)
            nxt = nxt.find_next_sibling()
        text = " ".join(texts)[:3000]
        if text: sections[matched_key or heading] = text
    return sections

# 读索引
with open(DATA_DIR / "drug_index.json") as f:
    drugs = json.load(f)

# 只取化学药+生物制品，跳过中药
candidates = [d for d in drugs if d["volume"] in (2, 3)]

out_dir = DATA_DIR / "wiki_drugs"
out_dir.mkdir(parents=True, exist_ok=True)
already = set(f.stem for f in out_dir.glob("*.json"))

total = 0; ok = 0
client = httpx.Client(timeout=20, follow_redirects=True)

print(f"🚀 启动采集: {len(candidates)} 候选, 已采 {len(already)}, {len(candidates)-len(already)} 待采\n")

for drug in candidates:
    name = drug["name"]
    safe = name.replace("/","_").replace(":","_")
    if safe in already: continue
    total += 1
    if total > 50: break  # 首轮限量50个
    try:
        r = client.get(f"https://zh.wikipedia.org/wiki/{quote(name)}", headers=H)
        if r.status_code != 200:
            print(f"  [{total}] ❌ {name} HTTP{r.status_code}")
            time.sleep(2); continue
        sections = parse_sections(r.text)
        if len(sections) <= 1:
            print(f"  [{total}] ⚠️ {name} (空)")
            time.sleep(2); continue
        entry = {"drug_id":f"WIKI-{name}","name":name,"sections":sections,
                 "source":{"version":"维基百科","volume":f"https://zh.wikipedia.org/wiki/{quote(name)}","page":"线上"}}
        with open(out_dir/f"{safe}.json","w") as f:
            json.dump(entry, f, ensure_ascii=False, indent=2)
        ok += 1
        print(f"  [{total}] ✅ {name} ({len(sections)}段)")
        time.sleep(3)
        if ok % 15 == 0 and ok > 0:
            print(f"  ⏸ 暂停45s...\n"); time.sleep(45)
    except Exception as e:
        print(f"  [{total}] ❌ {name}: {e}")
        time.sleep(1)

# 合并入库
merged = []
for f in out_dir.glob("*.json"):
    entry = json.load(open(f))
    if len(entry["sections"]) > 1:
        entry["category"] = "化学药"
        merged.append(entry)
Path(DATA_DIR / "wiki_merged.json").write_text(json.dumps(merged, ensure_ascii=False, indent=2))
print(f"\n✅ 完成: {ok} 新增 / {total} 处理, 合并 {len(merged)} 条 → wiki_merged.json")
client.close()