"""预生成所有文献的 AI 摘要（三种模式），存入 ai_summary JSON 列"""

import asyncio
import json

from sqlalchemy import select, text

from app.db import async_session
from app.models.literature import GlobalLiterature
from app.services.ai_summary import (
    generate_clinical_implication,
    generate_one_liner,
    generate_structured_summary,
)

MAX_CONCURRENT = 3


async def batch_generate(mode: str) -> int:
    gen_func = {
        "one_liner": generate_one_liner,
        "structured": generate_structured_summary,
        "implication": generate_clinical_implication,
    }[mode]

    async with async_session() as db:
        result = await db.execute(
            select(GlobalLiterature.id, GlobalLiterature.title, GlobalLiterature.abstract, GlobalLiterature.ai_summary)
            .where(GlobalLiterature.abstract.isnot(None), GlobalLiterature.abstract != "")
        )
        rows = result.all()

    todo = [(row.id, row.title, row.abstract, dict(row.ai_summary) if row.ai_summary else {})
            for row in rows if not (row.ai_summary or {}).get(mode)]
    total_need = len(todo)
    print(f"=== {mode}: {len(rows)} total, {total_need} need generation ===", flush=True)

    if total_need == 0:
        return 0

    sem = asyncio.Semaphore(MAX_CONCURRENT)

    async def call_ai(title: str, abstract: str) -> str | None:
        async with sem:
            return await gen_func(title, abstract)

    # Generate all summaries (no DB interaction needed)
    summaries = await asyncio.gather(*[call_ai(t, a) for _, t, a, _ in todo])

    # Build updates
    updates = []
    for (lit_id, _, _, curr), summary in zip(todo, summaries):
        if summary:
            curr[mode] = summary
            updates.append((lit_id, curr))

    # Write to DB via raw SQL in batches
    if updates:
        async with async_session() as db:
            BATCH = 100
            for i in range(0, len(updates), BATCH):
                batch = updates[i:i + BATCH]
                for lit_id, ai in batch:
                    await db.execute(
                        text("UPDATE global_literature SET ai_summary = CAST(:ai AS JSON), updated_at = NOW() WHERE id = :id"),
                        {"ai": json.dumps(ai), "id": lit_id},
                    )
                await db.commit()
                print(f"  {mode}: committed {min(i+BATCH, len(updates))}/{len(updates)}", flush=True)
    else:
        print(f"  {mode}: no summaries generated (AI service may be unavailable)")

    print(f"  {mode}: DONE — {len(updates)}/{total_need} generated", flush=True)
    return len(updates)


async def main():
    total = 0
    for mode in ("one_liner", "structured", "implication"):
        n = await batch_generate(mode)
        total += n
        print()
    print(f"All done! Generated {total} summaries total.", flush=True)


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