#!/usr/bin/env python3
"""PubMed FTP 基线导入 + 增量管道（生产脚本）"""

import argparse
import asyncio
import contextlib
import gzip
import hashlib
import os
import sys
import uuid
from datetime import date, datetime
from pathlib import Path

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

from app.compat import UTC

from lxml import etree
from sqlalchemy import and_, func, select

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 app.services.mesh_filter import get_mesh_filter
from app.services.tag_service import retag_all_unindexed


def uid(): 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 _extract_article(elem) -> dict | None:
    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
        title = article.findtext('.//ArticleTitle', '')
        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 = None
        for eid in article.findall('.//ELocationID'):
            if eid.get('EIdType') == 'doi': doi = eid.text; break
        pmc_id = None
        for aid in article.findall('.//ArticleIdList/ArticleId'):
            if aid.get('IdType') == 'pmc': pmc_id = aid.text; break
        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

        # CitationStatus（文献标引阶段）
        citation_status = medline.get('Status')

        # DateCompleted（MeSH 标引完成日）
        date_completed = None
        dc_elem = medline.find('.//DateCompleted')
        if dc_elem is not None:
            try:
                vr = int(dc_elem.findtext('Year', '0'))
                mr = int(dc_elem.findtext('Month', '1'))
                dr_ = int(dc_elem.findtext('Day', '1'))
                date_completed = datetime(vr, mr, dr_, tzinfo=UTC)
            except ValueError:
                pass

        # 临床试验注册号
        trial_reg = {}
        for aid in elem.findall('.//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

        # 撤稿追踪（CommentsCorrections）
        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

        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 ''})
        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, pages = 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:
                        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)
                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 ''
        pub_types = [pt.text for pt in article.findall('.//PublicationTypeList/PublicationType') if pt.text]
        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})

        # ── 补充解析字段 ──
        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

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

        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

        databank_list = []
        for db in medline.findall('.//DataBankList/DataBank'):
            db_name = db.findtext('DataBankName', '')
            accessions = [an.text for an in db.findall('.//AccessionNumberList/AccessionNumber') if an.text]
            if db_name:
                databank_list.append({'databank_name': db_name, 'accession_numbers': accessions})

        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.strip(), '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, 'keywords': keywords, 'grants': grants,
            'journal': journal_title, 'journal_issn': issn, 'journal_iso': journal_iso,
            'volume': volume, 'issue': issue, 'pages': pages or None,
            'pub_date': pub_date, 'pub_year': pub_year, 'pub_types': pub_types,
            'mesh_headings': mesh_headings, 'language': article.findtext('.//Language', 'en'),
            'chemical_list': chemical_list, 'gene_symbols': gene_symbols, 'num_refs': num_refs,
            'publication_status': publication_status, 'article_date': article_date,
            'databank_list': databank_list, 'suppl_mesh_list': suppl_mesh_list,
            'retracted': retracted, 'retraction_details': retraction_details,
        }
    except Exception:
        return None


def _is_oncology(article: dict) -> bool:
    """判断一篇文献是否属于肿瘤学范畴。

    两层过滤：
      Tier 1 — 有 MeSH 标引 → 仅检查 C04 树号（精确），不检查文本
                此规则基于数据分析：23.6 万条现有非 C04 误判文献中，
                有 MeSH 再检查文本会导致大量假阳性。
      Tier 2 — 无 MeSH 标引 → 关键词评分：
                - 标题命中任一关键词 → 收录（强信号）
                - 摘要 ≥2 关键词命中 → 收录
                - 摘要 0-1 关键词 → 放过（等待标引后回查）
    """
    mf = get_mesh_filter()
    mesh_headings = article.get('mesh_headings', [])

    if mesh_headings:
        # Tier 1: 有 MeSH → 只看 C04 树号
        return mf.is_oncology_by_mesh(mesh_headings)

    # Tier 2: 无 MeSH → 关键词评分模式
    return mf.is_oncology_by_text_score(
        article.get('title'), article.get('abstract')
    )


async def import_xml_file(xml_path: Path, stats: dict, min_year: int = 0, overwrite: bool = False) -> None:
    """导入单个 .xml.gz 文件"""
    count_new, count_skip, count_over = 0, 0, 0
    async with async_session() as db:
        with gzip.open(xml_path, 'rb') as f:
            try:
                context = etree.iterparse(f, events=('end',), tag='PubmedArticle')
            except Exception as e:
                print(f"  [ERROR] Failed to parse XML: {e}")
                return
            batch = []
            try:
                for _event, elem in context:
                    try:
                        article = _extract_article(elem)
                        elem.clear()
                        if not article: continue
                        if not _is_oncology(article):
                            count_skip += 1; continue
                        if min_year > 0 and (article.get('pub_year') or 0) < min_year:
                            count_skip += 1; continue
                        ex = await db.execute(select(GlobalLiterature).where(GlobalLiterature.pmid == article['pmid']))
                        existing = ex.scalar_one_or_none()
                        if existing:
                            if overwrite and article.get('pub_date'):
                                existing.pub_date = article['pub_date']
                                existing.pub_year = article['pub_year']
                                count_over += 1
                            else:
                                count_skip += 1
                            continue
                        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'),
                            volume=article.get('volume'), issue=article.get('issue'),
                            pages=article.get('pages'), pub_date=article.get('pub_date'),
                            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'),
                            grants=article.get('grants', []),
                            citation_status=article.get('citation_status'),
                            date_completed=article.get('date_completed'),
                            trial_reg=article.get('trial_reg'),
                            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); count_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 = []
                    except Exception:
                        continue
            except Exception as e:
                print(f"  [WARN] {xml_path.name}: XML iteration error at ~{count_new + count_skip}: {e}")
            if batch: await db.flush()
        await db.commit()
    stats['new'] += count_new; stats['skip'] += count_skip; stats['over'] += count_over; stats['files'] += 1
    label = f"+{count_new} new" + (f", {count_over} updated" if count_over else "")
    print(f"  {xml_path.name}: {label}, {count_skip} skipped ({stats['new']} total)")


async def validate_db() -> dict:
    """数据库完整性健康检查"""
    checks = {}
    async with async_session() as db:
        # 1. 总量
        total = await db.scalar(select(func.count(GlobalLiterature.id)))
        checks['total'] = total or 0

        # 2. PMID 范围
        pmid_min = await db.scalar(select(func.min(GlobalLiterature.pmid)))
        pmid_max = await db.scalar(select(func.max(GlobalLiterature.pmid)))
        checks['pmid_range'] = (pmid_min, pmid_max)

        # 3. 空字段统计
        null_title = await db.scalar(select(func.count(GlobalLiterature.id)).where(
            (GlobalLiterature.title == '') | GlobalLiterature.title.is_(None)))
        null_abstract = await db.scalar(select(func.count(GlobalLiterature.id)).where(
            GlobalLiterature.abstract.is_(None)))
        null_pub_year = await db.scalar(select(func.count(GlobalLiterature.id)).where(
            GlobalLiterature.pub_year.is_(None)))

        checks['null_title'] = null_title or 0
        checks['null_abstract'] = null_abstract or 0
        checks['null_pub_year'] = null_pub_year or 0

        # 4. 重复 PMID
        dup = await db.execute(
            select(GlobalLiterature.pmid, func.count(GlobalLiterature.id).label('cnt'))
            .group_by(GlobalLiterature.pmid)
            .having(func.count(GlobalLiterature.id) > 1)
        )
        dups = dup.all()
        checks['duplicate_pmids'] = len(dups)
        checks['duplicate_examples'] = [str(d[0]) for d in dups[:5]]

        # 5. 来源分布
        sources = await db.execute(
            select(GlobalLiterature.source, func.count(GlobalLiterature.id))
            .group_by(GlobalLiterature.source)
        )
        checks['sources'] = {s: c for s, c in sources.all()}

        # 6. retracted 标记
        retracted_count = await db.scalar(select(func.count(GlobalLiterature.id)).where(
            GlobalLiterature.retracted == True))
        checks['retracted'] = retracted_count or 0

        # 7. 年份分布概要（近5年 + 超老年份）
        recent = await db.execute(
            select(GlobalLiterature.pub_year, func.count(GlobalLiterature.id))
            .where(GlobalLiterature.pub_year >= 2020)
            .group_by(GlobalLiterature.pub_year)
            .order_by(GlobalLiterature.pub_year.desc())
        )
        checks['recent_years'] = {str(y): c for y, c in recent.all()}

        very_old = await db.scalar(select(func.count(GlobalLiterature.id)).where(
            GlobalLiterature.pub_year < 1900))
        checks['pre_1900'] = very_old or 0

        # 8. 标题过短（疑似损坏）
        short = await db.scalar(select(func.count(GlobalLiterature.id)).where(
            GlobalLiterature.title.isnot(None) & (func.length(GlobalLiterature.title) < 10)))
        checks['short_titles'] = short or 0

    # 健康状况判定
    issues = []
    if checks['null_title'] > 10000:
        issues.append(f"空标题过多: {checks['null_title']}")
    if checks['short_titles'] > 0:
        issues.append(f"短标题(疑似损坏): {checks['short_titles']}")
    if checks['duplicate_pmids'] > 0:
        issues.append(f"重复 PMID: {checks['duplicate_pmids']} 组")
    if checks['retracted'] == 0:
        issues.append("撤稿标记为 0（可能提取逻辑有遗漏）")
    if checks['pre_1900'] > 100:
        issues.append(f"1900 年前文献较多: {checks['pre_1900']}（正常应为个位数）")

    checks['healthy'] = len(issues) == 0
    checks['issues'] = issues
    return checks


async def main():
    parser = argparse.ArgumentParser(description='PubMed Baseline Importer')
    parser.add_argument('--dir', default=None, help='Path to PubMed baseline XML directory')
    parser.add_argument('--init-db', action='store_true', help='Create tables')
    parser.add_argument('--max-files', type=int, default=0, help='Max files to import (0=all)')
    parser.add_argument('--demo', action='store_true', help='Demo mode: only 5 files')
    parser.add_argument('--min-year', type=int, default=0, help='Minimum publication year (e.g. 2016)')
    parser.add_argument('--overwrite', action='store_true', help='Overwrite existing pub_date/pub_year for existing PMIDs')
    parser.add_argument('--validate', action='store_true', help='Run DB integrity check (can run standalone without --dir)')
    args = parser.parse_args()

    if args.validate:
        checks = await validate_db()
        print("=" * 50)
        print("  DB 完整性检查报告")
        print("=" * 50)
        print(f"  文献总量:          {checks['total']:,}")
        print(f"  PMID 范围:         {checks['pmid_range']}")
        print(f"  空标题:            {checks['null_title']:,}")
        print(f"  空摘要:            {checks['null_abstract']:,}")
        print(f"  无年份:            {checks['null_pub_year']:,}")
        print(f"  重复 PMID 组:      {checks['duplicate_pmids']}")
        if checks['duplicate_examples']:
            print(f"  重复样例:          {', '.join(checks['duplicate_examples'])}")
        print(f"  来源分布:          {checks['sources']}")
        print(f"  撤稿标记:          {checks['retracted']}")
        print(f"  1900年前文献:      {checks['pre_1900']:,}")
        print(f"  短标题(疑似损坏):  {checks['short_titles']:,}")
        print(f"  近5年分布:         {checks['recent_years']}")
        if checks['healthy']:
            print("  健康状况:          ✅ 正常")
        else:
            print("  健康状况:          ⚠️ 有异常")
            for iss in checks['issues']:
                print(f"    - {iss}")
        return

    if args.init_db:
        import subprocess
        script_dir = os.path.dirname(os.path.abspath(__file__))
        project_root = os.path.dirname(script_dir)
        subprocess.run(
            [sys.executable, "-m", "alembic", "-c", os.path.join("alembic", "alembic.ini"), "upgrade", "head"],
            cwd=project_root, check=True
        )
        print('Database tables created')

    if args.dir:
        files = sorted(Path(args.dir).glob('*.xml.gz'))
        limit = 5 if args.demo else (args.max_files if args.max_files > 0 else len(files))
        print(f'Found {len(files)} XML files (importing {limit})')

        stats = {'new': 0, 'skip': 0, 'over': 0, 'files': 0}
        for f in files[:limit]:
            await import_xml_file(f, stats, args.min_year, args.overwrite)

        async with async_session() as db:
            from app.models.operations import PipelineRun
            db.add(PipelineRun(id=uid(), run_type='baseline', status='success',
                    articles_new=stats['new'], articles_total=stats['new'] + stats['skip'],
                    files_processed=stats['files'],
                    started_at=datetime.now(UTC), completed_at=datetime.now(UTC)))
            await db.commit()
        print(f'Import complete: +{stats["new"]} new from {stats["files"]} files')

        # 标签回填（跳过，生产 GlobalTag 模型缺少新版字段）
        # print('Skipping retag_all_unindexed (model mismatch on prod)')

        if args.validate and stats['new'] > 0:
            print("\n" + "=" * 50)
            print("  导入后完整性验证")
            print("=" * 50)
            checks = await validate_db()
            print(f"  文献总量:          {checks['total']:,}")
            print(f"  空标题:            {checks['null_title']:,}")
            print(f"  空摘要:            {checks['null_abstract']:,}")
            print(f"  重复 PMID:         {checks['duplicate_pmids']}")
            if checks['healthy']:
                print("  验证结果:          ✅ 通过")
            else:
                print("  验证结果:          ⚠️ 有异常")
                for iss in checks['issues']:
                    print(f"    - {iss}")


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