"""
Full citation refresh - standalone script using direct SQL.
No app imports needed.
"""
import asyncio, sys
from datetime import datetime, timezone
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine

DB = 'postgresql+asyncpg://scilit:scilit_prod_2026@postgres:5432/scilit'
LOG = '/tmp/citation_refresh2.log'

ELINK_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/elink.fcgi"
BATCH = 200  # NCBI max per request
RATE = 3.0   # req/s
INTERVAL = 1.0 / RATE

def log(msg):
    ts = datetime.now(timezone.utc).isoformat()
    line = f'[{ts}] {msg}\n'
    print(line, end='', flush=True)
    with open(LOG, 'a') as f:
        f.write(line)

async def fetch_citedby(pmids):
    """Fetch citation counts for a batch of PMIDs via NCBI elink."""
    import httpx
    import xml.etree.ElementTree as ET

    pmid_str = ','.join(str(p) for p in pmids)
    url = f"{ELINK_URL}?dbfrom=pubmed&linkname=pubmed_pubmed_citedin&id={pmid_str}&retmode=xml"

    async with httpx.AsyncClient(timeout=30.0) as client:
        resp = await client.get(url)
        resp.raise_for_status()

    root = ET.fromstring(resp.content)
    result = {}
    for linkset in root.findall(".//LinkSet"):
        src_id_elem = linkset.find("IdList/Id")
        if src_id_elem is None or not src_id_elem.text:
            continue
        src_pmid = int(src_id_elem.text.strip())
        count = 0
        for lsdb in linkset.findall(".//LinkSetDb"):
            ln = lsdb.findtext("LinkName", "")
            if ln == "pubmed_pubmed_citedin":
                count = len([e for e in lsdb.findall(".//Link/Id") if e.text])
                break
        result[src_pmid] = max(result.get(src_pmid, 0), count)

    return result

async def main():
    log('Starting citation refresh (standalone)...')
    engine = create_async_engine(DB, pool_size=4)

    # Get all PMIDs
    async with engine.connect() as c:
        rows = (await c.execute(text("SELECT pmid FROM global_literature WHERE pmid IS NOT NULL ORDER BY created_at DESC"))).all()
    targets = [r[0] for r in rows]
    total = len(targets)
    log(f'Found {total} PMIDs')

    total_updated = 0
    total_errors = 0

    # Process ALL PMIDs via elink (the function handles its own batching)
    # We do it in chunks of 5000 for progress reporting and recovery
    CHUNK = 5000
    for i in range(0, total, CHUNK):
        chunk = targets[i:i+CHUNK]
        cn = i // CHUNK + 1
        tc = (total + CHUNK - 1) // CHUNK
        log(f'Chunk {cn}/{tc}: {len(chunk)} PMIDs ({i}/{total}, {i*100//total}%)')

        # Process each chunk's PMIDs via elink (internal 200/batch)
        all_counts = {}
        for j in range(0, len(chunk), BATCH):
            batch = chunk[j:j+BATCH]
            try:
                counts = await fetch_citedby(batch)
                all_counts.update(counts)
            except Exception as e:
                log(f'  elink sub-batch error: {e}')
            await asyncio.sleep(INTERVAL)

        # Write to DB
        updated = 0
        async with engine.connect() as c:
            for pmid, count in all_counts.items():
                await c.execute(
                    text("UPDATE global_literature SET cited_by_count = :c, updated_at = NOW() WHERE pmid = :p"),
                    {"c": count, "p": pmid},
                )
                updated += 1
            await c.commit()

        total_updated += updated
        total_errors += len(chunk) - len(all_counts)
        log(f'  => +{updated} updated ({total_updated}/{total})')

    log(f'DONE. Updated: {total_updated}, Errors: {total_errors}')
    await engine.dispose()

asyncio.run(main())