#!/usr/bin/env python3
"""Elasticsearch 索引重建脚本"""

import asyncio
import os
import sys

sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..'))

from sqlalchemy import select

from app.db import async_session
from app.models.literature import GlobalLiterature, GlobalLiteratureTag, GlobalTag
from app.services.search_service import search_service


async def rebuild_index(batch_size: int = 100):
    """全量重建 ES 索引"""
    async with async_session() as db:
        offset = 0
        total_indexed = 0
        while True:
            result = await db.execute(
                select(GlobalLiterature).order_by(GlobalLiterature.id).offset(offset).limit(batch_size)
            )
            articles = result.scalars().all()
            if not articles:
                break

            for lit in articles:
                # 加载标签
                tag_result = await db.execute(
                    select(GlobalTag.id, GlobalTag.name_zh, GlobalTag.path, GlobalTag.tag_category, GlobalLiteratureTag.is_major)
                    .join(GlobalLiteratureTag)
                    .where(GlobalLiteratureTag.literature_id == lit.id)
                )
                tags = [{"id": str(i), "name_zh": n, "path": p, "category": c, "is_major": m}
                        for i, n, p, c, m in tag_result]

                authors = lit.authors or []
                doc = {
                    "title": lit.title,
                    "abstract": lit.abstract or "",
                    "journal": lit.journal or "",
                    "pub_year": lit.pub_year,
                    "pub_date": lit.pub_date.isoformat() if lit.pub_date else None,
                    "authors": [{"family": a.get("family",""), "given": a.get("given","")} for a in authors],
                    "tag_ids": [t["id"] for t in tags],
                    "pmid": lit.pmid,
                    "doi": lit.doi or "",
                }
                await search_service.index_literature(str(lit.id), doc)
                total_indexed += 1

            offset += batch_size
            print(f"  Indexed {total_indexed} articles...")
            await asyncio.sleep(0.1)  # ES rate limit

    print(f"Index rebuild complete: {total_indexed} articles")


async def main():
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--batch-size", type=int, default=100)
    args = parser.parse_args()
    print("Rebuilding ES index...")
    await rebuild_index(args.batch_size)


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