"""从 mtrees2025.bin 导入 C04 MeSH Descriptor 到 global_tags + global_tag_tree_numbers

用法:
  # 统计模式（只预览，不下手）
  PYTHONPATH=. python scripts/import_mesh_tags.py --mode stats

  # 实际导入
  PYTHONPATH=. python scripts/import_mesh_tags.py --mode import

生产环境:
  # 先 docker cp mtrees 到容器
  docker cp /root/scilit/backend/app/data/mtrees2025.bin scilit-worker-1:/tmp/mtrees2025.bin
  docker cp backend/scripts/import_mesh_tags.py scilit-worker-1:/app/scripts/import_mesh_tags.py

  # 统计
  docker exec -e DATABASE_URL="postgresql+asyncpg://scilit:scilit_prod_2026@postgres:5432/scilit" \
    scilit-worker-1 python /app/scripts/import_mesh_tags.py --mode stats

  # 导入
  docker exec -e DATABASE_URL="postgresql+asyncpg://scilit:scilit_prod_2026@postgres:5432/scilit" \
    scilit-worker-1 python /app/scripts/import_mesh_tags.py --mode import
"""

import argparse
import asyncio
import os
import sys
from collections import defaultdict
from pathlib import Path

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
os.environ.setdefault("SPECIALTY", "oncology")

from sqlalchemy import select, func
from app.db import async_session
from app.models.literature import GlobalTag, GlobalTagTreeNumber

MTREES_PATH = Path(__file__).parent.parent / "app" / "data" / "mtrees2025.bin"


def parse_mtrees(path: str | Path) -> dict[str, set[str]]:
    """解析 mtrees 文件，返回 {descriptor_name: {tree_number, ...}}

    mtrees 格式: " DescriptorName;TreeNumber"（支持有/无 D 号列）
    只返回 C04 树号的记录。
    """
    desc_to_trees: dict[str, set[str]] = defaultdict(set)
    with open(path, encoding="utf-8") as f:
        for line in f:
            parts = line.strip().split(";")
            if len(parts) < 2:
                continue
            name = parts[0].strip()
            tree_number = parts[1].strip()
            if tree_number.startswith("C04"):
                desc_to_trees[name].add(tree_number)
    return dict(desc_to_trees)


def tree_category(tree_number: str) -> str:
    return tree_number.split(".")[0]


def tree_level(tree_number: str) -> int:
    return tree_number.count(".") + 1


async def scan_existing_tags(db) -> tuple[dict[str, GlobalTag], dict[str, GlobalTag]]:
    """返回 (name_en_lower → tag, mesh_ui → tag) 两个查找字典"""
    result = await db.execute(select(GlobalTag))
    tags = result.scalars().all()
    by_name: dict[str, GlobalTag] = {}
    by_mesh: dict[str, GlobalTag] = {}
    for t in tags:
        if t.name_en:
            by_name[t.name_en.lower()] = t
        if t.mesh_ui:
            by_mesh[t.mesh_ui] = t
    return by_name, by_mesh


async def stats_mode():
    desc_to_trees = parse_mtrees(MTREES_PATH)
    print(f"mtrees 中 C04 行数: {sum(len(v) for v in desc_to_trees.values())}")
    print(f"C04 唯一 Descriptor 数: {len(desc_to_trees)}")

    # 深度分布
    level_dist: dict[int, int] = defaultdict(int)
    for trees in desc_to_trees.values():
        for tn in trees:
            level_dist[tree_level(tn)] += 1
    print(f"\n树号深度分布:")
    for lv in sorted(level_dist):
        print(f"  深度 {lv}: {level_dist[lv]} 个树号")

    # 与现有标签的对比
    async with async_session() as db:
        by_name, by_mesh = await scan_existing_tags(db)

    matched = 0
    for name in desc_to_trees:
        if name.lower() in by_name:
            matched += 1
    print(f"\n现有 global_tags 中已匹配: {matched}/{len(desc_to_trees)}")
    print(f"将新建: {len(desc_to_trees) - matched}")

    # 样本
    print(f"\n前 5 个示例:")
    for i, (name, trees) in enumerate(sorted(desc_to_trees.items())[:5]):
        hit = "[match]" if name.lower() in by_name else "[new]"
        print(f"  {hit} {name} -> {', '.join(sorted(trees)[:3])}")


async def import_mode():
    desc_to_trees = parse_mtrees(MTREES_PATH)
    print(f"解析到 {len(desc_to_trees)} 个 C04 Descriptor")

    async with async_session() as db:
        by_name, by_mesh = await scan_existing_tags(db)
        print(f"当前 global_tags 共 {len(by_name)} 条")

        new_count = 0
        tree_count = 0
        skip_count = 0

        for name, trees in sorted(desc_to_trees.items()):
            name_lower = name.lower()
            existing = by_name.get(name_lower)

            if existing:
                # 已有标签 → 只补写 tree_numbers
                tag_id = existing.id
                skip_count += 1
            else:
                # 新建标签
                # 从 tree_numbers 推断浅层 tree_number（最少点号）
                primary_tn = min(trees, key=lambda tn: tree_level(tn))
                tag = GlobalTag(
                    name_zh=None,   # C04 无中文名，前台显示 name_en
                    name_en=name,
                    path=name,
                    level=tree_level(primary_tn),
                    tag_category="cancer",
                    is_selectable=tree_level(primary_tn) <= 3,
                    source="mesh",
                    is_active=True,
                )
                db.add(tag)
                await db.flush()
                tag_id = tag.id
                by_name[name_lower] = tag  # 防止同一批内重复创建
                new_count += 1

            # 写入 tree_numbers
            for tn in trees:
                # 检查是否已有这条 tree_number（防止重复导入）
                existing_tn = await db.execute(
                    select(GlobalTagTreeNumber).where(
                        GlobalTagTreeNumber.tag_id == tag_id,
                        GlobalTagTreeNumber.tree_number == tn,
                    )
                )
                if existing_tn.scalar():
                    continue
                is_primary = (tn == min(trees, key=lambda x: tree_level(x)))
                db.add(GlobalTagTreeNumber(
                    tag_id=tag_id,
                    tree_number=tn,
                    tree_category="C04",
                    is_primary=is_primary,
                ))
                tree_count += 1

            if (new_count + skip_count) % 500 == 0:
                print(f"  进度: {new_count} 新建, {skip_count} 已有, {tree_count} 树号")
                await db.commit()

        await db.commit()

    total_tags = await asyncio.to_thread(
        lambda: sum(1 for _ in desc_to_trees)
    )
    print(f"\n完成! 新建 {new_count} 个标签, 跳过 {skip_count} 个已有, 写入 {tree_count} 个树号")


async def main():
    parser = argparse.ArgumentParser(description="导入 C04 MeSH 标签到 global_tags")
    parser.add_argument("--mode", choices=["stats", "import"], default="stats",
                        help="stats: 只统计; import: 实际导入")
    args = parser.parse_args()

    if not MTREES_PATH.exists():
        print(f"错误: mtrees 文件不存在: {MTREES_PATH}")
        sys.exit(1)

    if args.mode == "stats":
        await stats_mode()
    else:
        await import_mode()


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