"""研究设计自动分类服务：在管道打标后运行，将 pub_types 映射为分层分类"""

import uuid

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

from app.constants.study_design import classify_study_design
from app.models.literature import GlobalLiterature


async def classify_literature_design(db: AsyncSession, lit_id: uuid.UUID) -> dict | None:
    """为单篇文献分类研究设计，写回 study_design 字段"""
    lit = (await db.execute(select(GlobalLiterature).where(GlobalLiterature.id == lit_id))).scalar()
    if not lit or not lit.pub_types:
        return None
    result = classify_study_design(lit.pub_types)
    if result and lit.study_design != result:
        lit.study_design = result
    return result


async def batch_classify_designs(
    db: AsyncSession,
    lit_ids: list[uuid.UUID],
) -> int:
    """批量分类，返回更新的文献数"""
    if not lit_ids:
        return 0
    result = await db.execute(
        select(GlobalLiterature).where(GlobalLiterature.id.in_(lit_ids))
    )
    updated = 0
    for lit in result.scalars().all():
        if not lit.pub_types:
            continue
        design = classify_study_design(lit.pub_types)
        if design and lit.study_design != design:
            lit.study_design = design
            updated += 1
    if updated:
        await db.commit()
    return updated


async def backfill_all_designs(db: AsyncSession, batch_size: int = 200) -> int:
    """全库回填：为无 study_design 且有 pub_types 的文献计算分类"""
    total = 0
    while True:
        result = await db.execute(
            select(GlobalLiterature.id, GlobalLiterature.pub_types)
            .where(GlobalLiterature.study_design.is_(None), GlobalLiterature.pub_types.isnot(None))
            .limit(batch_size)
        )
        rows = result.all()
        if not rows:
            break
        updated = 0
        for lit_id, pub_types in rows:
            design = classify_study_design(pub_types or [])
            if design:
                await db.execute(
                    GlobalLiterature.__table__.update()
                    .where(GlobalLiterature.id == lit_id)
                    .values(study_design=design)
                )
                updated += 1
        if updated:
            await db.commit()
        total += updated
    return total
