#!/usr/bin/env python3
"""清空 global_literature 并用单篇基线 XML 文件导入真实文献数据。"""
import asyncio
import hashlib
import sys
import uuid
from datetime import date, datetime
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))

from lxml import etree
from sqlalchemy import text

from app.compat import UTC
from app.constants.study_design import classify_study_design
from app.db import async_session
from app.models.literature import GlobalLiterature
from app.services.journal_utils import ensure_journal_async
from scripts.pubmed_baseline import _is_oncology


def uid() -> uuid.UUID:
    return uuid.uuid4()


def _month_num(raw: str) -> int:
    months = {"jan": 1, "feb": 2, "mar": 3, "apr": 4, "may": 5, "jun": 6,
              "jul": 7, "aug": 8, "sep": 9, "oct": 10, "nov": 11, "dec": 12}
    try:
        return int(raw)
    except ValueError:
        return months.get(raw.lower()[:3], 1)


def _month_day(year: int, month: int, raw_day: str) -> int:
    try:
        day = int(raw_day)
        max_days = {1: 31, 2: 29 if year % 4 == 0 and (year % 100 != 0 or year % 400 == 0) else 28,
                    3: 31, 4: 30, 5: 31, 6: 30, 7: 31, 8: 31, 9: 30, 10: 31, 11: 30, 12: 31}
        return min(day, max_days.get(month, 31))
    except ValueError:
        return 1


def _parse_date_ymd(elem) -> datetime | None:
    """解析 PubmedData/History 下的 Year/Month/Day → 带 UTC 的 datetime"""
    if elem is None:
        return None
    try:
        y = int(elem.findtext("Year", "0"))
        m = int(elem.findtext("Month", "1"))
        d = int(elem.findtext("Day", "1"))
        return datetime(y, m, d, tzinfo=UTC)
    except (ValueError, TypeError):
        return None


def extract_article(elem) -> dict | None:
    """解析单条 PubmedArticle，含全部 6 个 PubMed 日期字段。"""
    try:
        medline = elem.find(".//MedlineCitation")
        article = medline.find(".//Article") if medline is not None else None
        if article is None:
            return None
        pmid_elem = medline.find(".//PMID")
        pmid = int(pmid_elem.text) if pmid_elem is not None and pmid_elem.text else None
        if not pmid:
            return None

        # ── 标题（itertext 捕获 <i> 等子标签）──
        title_elem = article.find(".//ArticleTitle")
        title = "".join(title_elem.itertext()).strip() if title_elem is not None else ""

        # ── 摘要 ──
        abstract_parts = []
        for at in article.findall(".//Abstract/AbstractText"):
            label = at.get("Label", "")
            text = "".join(at.itertext())
            abstract_parts.append(f"{label}: {text}" if label else text)
        abstract = "\n".join(abstract_parts) if abstract_parts else None

        # ── 作者 ──
        authors = [
            {"family": au.findtext("LastName", ""), "given": au.findtext("ForeName", ""),
             "affiliation": (au.findtext(".//AffiliationInfo/Affiliation") or "")}
            for au in article.findall(".//Author")
        ]

        # ── DOI ──
        doi = None
        for eid in article.findall(".//ELocationID"):
            if eid.get("EIdType") == "doi":
                doi = eid.text
                break

        # ── PMC ID（统一去掉 "PMC" 前缀）──
        pmc_id = None
        for aid in elem.findall(".//PubmedData/ArticleIdList/ArticleId"):  # PubmedData 层级
            if aid.get("IdType") == "pmc":
                raw = (aid.text or "").lstrip("PMC")
                pmc_id = raw if raw else None
                break
        if not pmc_id:
            for aid in article.findall(".//ArticleIdList/ArticleId"):  # 回退 Article 层级
                if aid.get("IdType") == "pmc":
                    raw = (aid.text or "").lstrip("PMC")
                    pmc_id = raw if raw else None
                    break

        # ── PubMed 修订日期 ──
        pubmed_revised = None
        dr_elem = medline.find(".//DateRevised")
        if dr_elem is not None:
            try:
                vr = int(dr_elem.findtext("Year", "0"))
                mr = int(dr_elem.findtext("Month", "1"))
                dr_ = int(dr_elem.findtext("Day", "1"))
                pubmed_revised = datetime(vr, mr, dr_, tzinfo=UTC)
            except ValueError:
                pass

        # ── 其他 MedlineCitation 元数据 ──
        citation_status = medline.get("Status")
        date_completed = _parse_date_ymd(medline.find(".//DateCompleted"))

        # ── 临床试验注册号 ──
        trial_reg = {}
        for aid in elem.findall(".//PubmedData/ArticleIdList/ArticleId"):
            id_type = aid.get("IdType")
            if id_type == "ClinicalTrials.gov":
                trial_reg["nct"] = aid.text
            elif id_type == "EU-CTR":
                trial_reg["eudract"] = aid.text
            elif id_type == "ChiCTR":
                trial_reg["chictr"] = aid.text

        # ── 撤稿 ──
        retracted = False
        retraction_details = None
        for cc in medline.findall(".//CommentsCorrections"):
            ref_type = cc.findtext("RefType")
            ref_pmid_elem = cc.find("PMID")
            ref_pmid = int(ref_pmid_elem.text) if ref_pmid_elem is not None and ref_pmid_elem.text else None
            if ref_type in ("RetractedBy", "RetractionOf"):
                retraction_details = {
                    "ref_type": ref_type,
                    "ref_pmid": ref_pmid,
                    "ref_source": cc.findtext("RefSource"),
                }
                if ref_type == "RetractedBy":
                    retracted = True

        # ── 期刊 ──
        j = article.find(".//Journal")
        journal_title = j.findtext("Title") if j is not None else None
        journal_iso = j.findtext("ISOAbbreviation") if j is not None else None
        issn = (j.findtext(".//ISSN") or "") if j is not None else None

        # ── 卷期页码 ──
        vi = j.find(".//JournalIssue") if j is not None else None
        volume, issue, pub_date, pub_year, print_date, pages = None, None, None, None, None, ""
        if vi is not None:
            volume = vi.findtext("Volume")
            issue = vi.findtext("Issue")
            pd = vi.find(".//PubDate")
            if pd is not None:
                y = pd.findtext("Year")
                if y:
                    pub_year = int(y)
                    m = pd.findtext("Month", "1")
                    d = pd.findtext("Day", "1")
                    try:
                        print_date = date(int(y), _month_num(m), _month_day(int(y), _month_num(m), d))
                    except ValueError:
                        print_date = date(int(y), 1, 1)
                else:
                    md = pd.findtext("MedlineDate")
                    if md:
                        import re
                        ym = re.search(r"(\d{4})", md)
                        if ym:
                            pub_year = int(ym.group(1))
            pages = article.findtext(".//Pagination/MedlinePgn") or ""

        # ── PubmedData/PublicationStatus ──
        publication_status = None
        ps_elem = elem.find(".//PubmedData/PublicationStatus")
        if ps_elem is not None:
            publication_status = ps_elem.text

        # ── [EPDAT] ArticleDate（电子出版日期）──
        article_date = None
        ad_elem = article.find(".//ArticleDate")
        if ad_elem is not None:
            try:
                ayr = int(ad_elem.findtext("Year", "0"))
                amo = int(ad_elem.findtext("Month", "1"))
                ady = int(ad_elem.findtext("Day", "1"))
                article_date = date(ayr, amo, ady)
            except ValueError:
                pass

        # ── [DP] = LEAST(article_date, print_date) ──
        effective_pub_date = None
        if article_date and print_date:
            effective_pub_date = min(article_date, print_date)
        elif article_date:
            effective_pub_date = article_date
        elif print_date:
            effective_pub_date = print_date

        # ── [CRDT] / [EDAT] PubmedData/History ──
        create_date = None   # CRDT — PubStatus="pubmed"
        entrez_date = None   # EDAT — PubStatus="entrez"
        for hist in elem.findall(".//PubmedData/History/PubMedPubDate"):
            status = hist.get("PubStatus", "")
            h_dt = _parse_date_ymd(hist)
            if h_dt:
                if status == "pubmed":
                    create_date = h_dt
                elif status == "entrez":
                    entrez_date = h_dt

        # ── 文献类型 ──
        pub_types = [pt.text for pt in article.findall(".//PublicationTypeList/PublicationType") if pt.text]

        # ── MeSH ──
        mesh_headings = []
        for mh in medline.findall(".//MeshHeadingList/MeshHeading"):
            desc = mh.findtext("DescriptorName")
            dn = mh.find("DescriptorName")
            ui = dn.get("UI", "") if dn is not None else ""
            major = dn.get("MajorTopicYN", "N") == "Y" if dn is not None else False
            mesh_headings.append({"descriptor": desc, "ui": ui, "major": major})

        # ── 关键词 ──
        keywords = [kw.text for kw in article.findall(".//KeywordList/Keyword") if kw.text]

        # ── 基金 ──
        grants = []
        for gr in article.findall(".//GrantList/Grant"):
            grants.append({"abbreviation": gr.findtext("Acronym") or "",
                          "agency": gr.findtext("Agency") or "",
                          "country": gr.get("Country", ""),
                          "grant_id": gr.findtext("GrantID") or ""})

        # ── 化学物质 ──
        chemical_list = []
        for ch in medline.findall(".//ChemicalList/Chemical"):
            rn = ch.findtext("RegistryNumber", "")
            ns = ch.find("NameOfSubstance")
            name = ns.text if ns is not None else ""
            ui = ns.get("UI", "") if ns is not None else ""
            chemical_list.append({"name": name, "registry_number": rn, "mesh_ui": ui})

        # ── 基因 ──
        gene_symbols = [gs.text for gs in article.findall(".//GeneSymbolList/GeneSymbol") if gs.text]

        # ── 引用数 ──
        num_refs = None
        nr_elem = medline.find(".//NumberOfReferences")
        if nr_elem is not None and nr_elem.text:
            try:
                num_refs = int(nr_elem.text)
            except ValueError:
                pass

        # ── 数据库来源 ──
        databank_list = []
        for db_elem in medline.findall(".//DataBankList/DataBank"):
            db_name = db_elem.findtext("DataBankName", "")
            accessions = [an.text for an in db_elem.findall(".//AccessionNumberList/AccessionNumber") if an.text]
            if db_name:
                databank_list.append({"databank_name": db_name, "accession_numbers": accessions})

        # ── 补充 MeSH ──
        suppl_mesh_list = []
        for sm in medline.findall(".//SupplMeshList/SupplMeshName"):
            name = sm.text or ""
            ui = sm.get("UI", "")
            stype = sm.get("Type", "")
            suppl_mesh_list.append({"descriptor": name, "ui": ui, "type": stype})

        return {
            "pmid": pmid,
            "title": title,
            "abstract": abstract.strip() if abstract else None,
            "authors": authors,
            "doi": doi,
            "pmc_id": pmc_id,
            "is_oa": pmc_id is not None,
            "pubmed_revised": pubmed_revised,
            "citation_status": citation_status,
            "date_completed": date_completed,
            "trial_reg": trial_reg,
            "retracted": retracted,
            "retraction_details": retraction_details,
            "journal": journal_title,
            "journal_issn": issn,
            "journal_iso": journal_iso,
            "volume": volume,
            "issue": issue,
            "pages": pages or None,
            "print_date": print_date,                  # [PPDAT]
            "article_date": article_date,              # [EPDAT]
            "pub_date": effective_pub_date,            # [DP] = LEAST(article_date, print_date)
            "pub_year": pub_year,
            "publication_status": publication_status,
            "create_date": create_date,                # [CRDT]
            "entrez_date": entrez_date,                # [EDAT]
            "pub_types": pub_types,
            "mesh_headings": mesh_headings,
            "language": article.findtext(".//Language", "en"),
            "keywords": keywords,
            "grants": grants,
            "chemical_list": chemical_list,
            "gene_symbols": gene_symbols,
            "num_refs": num_refs,
            "databank_list": databank_list,
            "suppl_mesh_list": suppl_mesh_list,
        }
    except Exception:
        return None


async def import_baseline_file(xml_path: Path) -> dict:
    """清空并导入单篇基线 XML 文件（非 gzip）。返回统计。"""
    stats = {"new": 0, "skipped": 0, "total_articles": 0}

    # 1. 清空数据表
    async with async_session() as db:
        print("清空 global_literature（CASCADE 自动清除关联表）...")
        await db.execute(text("TRUNCATE TABLE global_literature CASCADE"))
        await db.commit()
        print("已清空。")

    # 2. 流式解析 XML
    total_articles = 0
    batch = []
    async with async_session() as db:
        context = etree.iterparse(str(xml_path), events=("end",), tag="PubmedArticle")
        for _event, elem in context:
            total_articles += 1
            article = extract_article(elem)
            elem.clear()
            # 清理父节点释放内存
            parent = elem.getparent()
            if parent is not None:
                parent.remove(elem)

            if not article:
                stats["skipped"] += 1
                continue

            # Oncology 过滤
            if not _is_oncology(article):
                stats["skipped"] += 1
                continue

            # 构建 GlobalLiterature 记录（全部字段）
            lit = GlobalLiterature(
                id=uid(),
                pmid=article["pmid"],
                title=article["title"],
                abstract=article.get("abstract"),
                authors=article.get("authors", []),
                doi=article.get("doi"),
                journal=article.get("journal"),
                journal_issn=article.get("journal_issn"),
                journal_iso=article.get("journal_iso"),
                volume=article.get("volume"),
                issue=article.get("issue"),
                pages=article.get("pages"),
                pub_date=article.get("pub_date"),            # [DP]
                print_date=article.get("print_date"),        # [PPDAT]
                pub_year=article.get("pub_year"),
                pub_types=article.get("pub_types", []),
                mesh_headings=article.get("mesh_headings", []),
                language=article.get("language", "en"),
                pmc_id=article.get("pmc_id"),
                is_oa=article.get("is_oa", False),
                keywords=article.get("keywords", []),
                pubmed_revised=article.get("pubmed_revised"),
                citation_status=article.get("citation_status"),
                date_completed=article.get("date_completed"),
                article_date=article.get("article_date"),    # [EPDAT]
                create_date=article.get("create_date"),      # [CRDT]
                entrez_date=article.get("entrez_date"),      # [EDAT]
                grants=article.get("grants", []),
                chemical_list=article.get("chemical_list", []),
                gene_symbols=article.get("gene_symbols", []),
                num_refs=article.get("num_refs"),
                publication_status=article.get("publication_status"),
                trial_reg=article.get("trial_reg"),
                databank_list=article.get("databank_list", []),
                suppl_mesh_list=article.get("suppl_mesh_list", []),
                study_design=classify_study_design(article.get("pub_types", [])),
                retracted=article.get("retracted", False),
                retraction_details=article.get("retraction_details"),
                source="pubmed_ftp",
                raw_xml_hash=hashlib.sha256(article["title"].encode()).hexdigest()[:16],
            )
            db.add(lit)
            batch.append(lit)
            stats["new"] += 1

            try:
                await ensure_journal_async(db, article.get("journal_issn"), article.get("journal"))
            except Exception:
                pass

            if len(batch) >= 100:
                await db.flush()
                batch = []

        if batch:
            await db.flush()
        await db.commit()

    stats["total_articles"] = total_articles
    return stats


async def main():
    # 真实文件在 D:/Lit/Baseline/pubmed26n0001.xml/pubmed26n0001.xml（目录同名）
    xml_path = Path("D:/Lit/Baseline/pubmed26n0001.xml/pubmed26n0001.xml")
    if not xml_path.exists():
        # 回退：gzip 版本
        xml_path = Path("D:/Lit/Baseline/pubmed26n0001.xml.gz")
        if not xml_path.exists():
            print(f"文件不存在: {xml_path}")
            sys.exit(1)

    file_size = xml_path.stat().st_size / 1024 / 1024
    print(f"开始导入: {xml_path.name} ({file_size:.0f} MB)")
    stats = await import_baseline_file(xml_path)

    print(f"\n导入完成:")
    print(f"  文件总文献数: {stats['total_articles']}")
    print(f"  新增(肿瘤):  {stats['new']}")
    print(f"  跳过(非肿瘤): {stats['skipped']}")

    # 验证
    async with async_session() as db:
        from sqlalchemy import select, func
        total = await db.scalar(select(func.count()).select_from(GlobalLiterature))
        print(f"\n  数据库中:    {total} 篇肿瘤文献")

        # 抽样检查新字段
        from sqlalchemy import select as sa_select
        row = (await db.execute(
            sa_select(GlobalLiterature.pmid, GlobalLiterature.journal_iso,
                      GlobalLiterature.print_date, GlobalLiterature.create_date,
                      GlobalLiterature.entrez_date, GlobalLiterature.article_date,
                      GlobalLiterature.pub_date)
            .limit(3)
        )).fetchall()
        for r in row:
            print(f"  PMID={r.pmid} iso={r.journal_iso} ppdat={r.print_date}"
                  f" crdt={str(r.create_date)[:19] if r.create_date else None}"
                  f" edat={str(r.entrez_date)[:19] if r.entrez_date else None}"
                  f" epdat={r.article_date} dp={r.pub_date}")


if __name__ == "__main__":
    asyncio.run(main())
