"""PubMed 数据管道：FTP 下载 → XML 流式解析 → MeSH 过滤 → 入库"""

from __future__ import annotations

import gzip
import hashlib
import re
from collections.abc import AsyncIterator
from datetime import date, datetime

from lxml import etree
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.compat import UTC
from app.db import async_session
from app.models.literature import GlobalLiterature
from app.services.journal_utils import ensure_journal_async
from app.services.tag_service import tag_article

# DOI 正则：从 PDF 全文提取
DOI_PATTERN = re.compile(r'10\.\d{4,9}/[-._;()/:A-Za-z0-9]+')

async def parse_xml_stream(xml_path: str) -> AsyncIterator[dict]:
    """流式解析 PubMed XML，逐条 yield 文献 dict"""
    with gzip.open(xml_path, "rb") as f:
        context = etree.iterparse(f, events=("end",), tag="PubmedArticle")

        for _event, elem in context:
            article_data = extract_article(elem)
            # 立即释放内存
            elem.clear()
            while elem.getprevious() is not None:
                del elem.getparent()[0]

            if article_data:
                yield article_data


def extract_article(elem) -> dict | None:
    """从单个 PubmedArticle XML 元素提取结构化数据"""
    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 pmid is None:
            return None

        title_elem = article.find(".//ArticleTitle")
        title = title_elem.text or "" if title_elem is not None else ""

        abstract_parts = []
        if article.find(".//Abstract/AbstractText") is not None:
            for at in article.findall(".//Abstract/AbstractText"):
                label = at.get("Label", "")
                text = at.text or ""
                abstract_parts.append(f"{label}: {text}" if label else text)
        abstract = "\n".join(abstract_parts) if abstract_parts else None

        authors = []
        for au in article.findall(".//Author"):
            last = au.findtext("LastName", "")
            fore = au.findtext("ForeName", "")
            aff = au.findtext(".//AffiliationInfo/Affiliation", "")
            authors.append({"family": last, "given": fore, "affiliation": aff})

        doi_elem = article.find(".//ELocationID[@EIdType='doi']")
        doi = doi_elem.text if doi_elem is not None else None

        journal_elem = article.find(".//Journal")
        journal_title = journal_elem.findtext("Title") if journal_elem is not None else None
        journal_iso = journal_elem.findtext("ISOAbbreviation") if journal_elem is not None else None
        issn_elem = journal_elem.find(".//ISSN") if journal_elem is not None else None
        journal_issn = issn_elem.text if issn_elem is not None else None

        vi = journal_elem.find(".//JournalIssue") if journal_elem is not None else None
        volume = vi.findtext("Volume") if vi is not None else None
        issue = vi.findtext("Issue") if vi is not None else None

        pub_date = None
        pub_year = None
        if vi is not None:
            pd_elem = vi.find(".//PubDate")
            if pd_elem is not None:
                y = pd_elem.findtext("Year")
                m = pd_elem.findtext("Month", "01")
                d = pd_elem.findtext("Day", "01")
                if y:
                    pub_year = int(y)
                    try:
                        pub_date = date(int(y), _month_num(m), _month_day(int(y), _month_num(m), d))
                    except ValueError:
                        pub_date = date(int(y), 1, 1)

        pagination = article.findtext(".//Pagination/MedlinePgn", "")

        pub_types = [pt.text for pt in article.findall(".//PublicationTypeList/PublicationType") if pt.text]

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

        lang_elem = article.find(".//Language")
        language = lang_elem.text if lang_elem is not None else "en"

        # 计算原始 XML 哈希用于去重
        raw_xml = etree.tostring(elem, encoding="unicode")
        raw_hash = hashlib.sha256(raw_xml.encode()).hexdigest()

        return {
            "pmid": pmid,
            "title": title,
            "abstract": abstract,
            "authors": authors,
            "doi": doi,
            "journal": journal_title,
            "journal_issn": journal_issn,
            "journal_iso": journal_iso,
            "volume": volume,
            "issue": issue,
            "pages": pagination or None,
            "pub_date": pub_date,
            "pub_year": pub_year,
            "pub_types": pub_types,
            "mesh_headings": mesh_headings,
            "language": language,
            "raw_xml_hash": raw_hash,
        }
    except Exception:
        return None


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


async def upsert_article(article_data: dict, db: AsyncSession | None = None) -> bool:
    """插入或更新单篇文献，返回 True 表示新插入

    Args:
        article_data: 文献数据字典
        db: 可选 — 传入外部 session（如测试用），由调用方管理 commit/close
    """
    if db is not None:
        return await _upsert_article_impl(article_data, db)
    async with async_session() as s:
        return await _upsert_article_impl(article_data, s)


async def _upsert_article_impl(article_data: dict, db: AsyncSession) -> bool:
    existing = await db.execute(
        select(GlobalLiterature.id).where(GlobalLiterature.pmid == article_data["pmid"])
    )
    lit_id = existing.scalar()

    if lit_id:
        result = await db.execute(select(GlobalLiterature).where(GlobalLiterature.id == lit_id))
        lit = result.scalar()
        if lit:
            lit.title = article_data["title"]
            lit.abstract = article_data["abstract"]
            lit.authors = article_data["authors"]
            lit.doi = article_data["doi"]
            lit.mesh_headings = article_data.get("mesh_headings", lit.mesh_headings)
            lit.updated_at = datetime.now(UTC)
        mh = article_data.get("mesh_headings", [])
        if mh:
            try:
                await tag_article(db, lit.id, mh)
            except Exception:
                pass
        await db.commit()
        return False
    else:
        lit = GlobalLiterature(**article_data)
        db.add(lit)
        try:
            await ensure_journal_async(db, article_data.get("journal_issn"), article_data.get("journal"))
        except Exception:
            pass
        # 需要 flush 才能拿到 lit.id
        try:
            await db.flush()
            mh = article_data.get("mesh_headings", [])
            if mh:
                await tag_article(db, lit.id, mh)
        except Exception:
            pass
        await db.commit()
        return True