"""Run FULL citation refresh with incremental progress logging"""
import asyncio, sys, os
from datetime import datetime
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker
from sqlalchemy import select
from app.models.literature import GlobalLiterature
from app.services.pubmed_api import fetch_citedby_counts, update_citation_counts

DATABASE_URL = 'postgresql+asyncpg://scilit:scilit_prod_2026@postgres:5432/scilit'

async def main():
    engine = create_async_engine(DATABASE_URL, pool_size=4)
    async_session = sessionmaker(engine, class_=AsyncSession)

    print(f"[{datetime.now().isoformat()}] Starting full citation refresh...", flush=True)

    async with async_session() as db:
        # Get all PMIDs
        result = await db.execute(select(GlobalLiterature.pmid).order_by(GlobalLiterature.created_at.desc()))
        all_pmids = [r for (r,) in result]
        total = len(all_pmids)
        print(f"[{datetime.now().isoformat()}] Found {total} PMIDs to process", flush=True)

        # Process in chunks of 5000 (matching the admin API limit=5000 behavior)
        BATCH = 5000
        total_updated = 0
        total_errors = 0

        for i in range(0, len(all_pmids), BATCH):
            batch = all_pmids[i:i+BATCH]
            counts = await fetch_citedby_counts(batch)

            async with async_session()() as s:
                updated = 0
                for pmid, cited in counts.items():
                    from sqlalchemy import update as sql_update
                    await s.execute(
                        sql_update(GlobalLiterature)
                        .where(GlobalLiterature.pmid == pmid)
                        .values(cited_by_count=cited, updated_at=datetime.now())
                    )
                    updated += 1
                await s.commit()

            total_updated += updated
            errors_in_batch = len(batch) - len(counts)
            total_errors += errors_in_batch
            pct = min(100, (i + BATCH) / total * 100)
            print(f"[{datetime.now().isoformat()}] batch {i//BATCH+1}/{(total+BATCH-1)//BATCH}: +{updated} updated ({pct:.0f}%)", flush=True)

    print(f"[{datetime.now().isoformat()}] Done. Updated: {total_updated}, Errors: {total_errors}", flush=True)
    await engine.dispose()

asyncio.run(main())