#!/usr/bin/env python3
"""从 PubMed 公开 E-utilities 取结直肠癌腹膜转移文献的 MeSH 标注。"""
import json, time, urllib.parse, urllib.request, pathlib
KEY = pathlib.Path.home().joinpath(".config/ncbi/api_key").read_text().strip()
EU = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
TERM = ('("colorectal neoplasms"[MeSH] OR colorectal cancer[Title/Abstract]) AND '
        '("peritoneal neoplasms"[MeSH] OR peritoneal metastas*[Title/Abstract] OR '
        'peritoneal carcinomatosis[Title/Abstract]) AND 2015:2026[dp]')


def get(path, **kw):
    kw["api_key"] = KEY
    with urllib.request.urlopen(f"{EU}/{path}?{urllib.parse.urlencode(kw)}", timeout=180) as fh:
        return fh.read()


ids = json.loads(get("esearch.fcgi", db="pubmed", term=TERM, retmax=10000,
                     retmode="json"))["esearchresult"]["idlist"]
print(f"  检索命中 {len(ids)} 篇")
recs = []
for i in range(0, len(ids), 200):
    xml = get("efetch.fcgi", db="pubmed", id=",".join(ids[i:i + 200]), retmode="xml").decode("utf-8", "ignore")
    for art in xml.split("<PubmedArticle>")[1:]:
        yr = None
        for tag in ("<Year>", "<MedlineDate>"):
            if tag in art:
                seg = art.split(tag, 1)[1][:12]
                digits = "".join(ch for ch in seg if ch.isdigit())[:4]
                if len(digits) == 4:
                    yr = int(digits); break
        mesh = []
        for chunk in art.split("<DescriptorName")[1:]:
            if ">" in chunk:
                mesh.append(chunk.split(">", 1)[1].split("<")[0].strip())
        if yr and mesh:
            recs.append({"year": yr, "mesh": sorted(set(mesh))})
    print(f"  已解析 {len(recs)} 篇", end="\r")
    time.sleep(0.15)
json.dump({"term": TERM, "n": len(recs), "records": recs}, open("corpus.json", "w"), ensure_ascii=False)
print(f"\n  有年份且有 MeSH 的: {len(recs)} 篇")