"""
Fix pub_date on records with YYYY-01-01 by querying NCBI E-utilities efetch.

Batch queries existing PMIDs via the efetch API, parses real dates from XML,
and updates the database. Much faster than re-running the full FTP baseline import
(~5 min for 147k PMIDs at 3 req/s).

Usage:
    python scripts/fix_dates_via_efetch.py [--limit N] [--batch 200] [--rate 3]

Requires DATABASE_URL env var or uses asyncdb default.
"""

import asyncio
import os
import sys
import time
import uuid
from datetime import date, datetime
from typing import Optional, Dict

import httpx
import lxml.etree as ET
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.orm import sessionmaker

from dateutil.parser import parse as parse_date

DATABASE_URL = os.environ.get(
    "DATABASE_URL",
    "postgresql+asyncpg://scilit:scilit@localhost:5432/scilit",
)

# ── helpers (mirrored from pubmed_api.py) ──

MONTH_MAP: Dict[str, int] = {
    "jan": 1, "january": 1, "1": 1,
    "feb": 2, "february": 2, "2": 2,
    "mar": 3, "march": 3, "3": 3,
    "apr": 4, "april": 4, "4": 4,
    "may": 5, "5": 5,
    "jun": 6, "june": 6, "6": 6,
    "jul": 7, "july": 7, "7": 7,
    "aug": 8, "august": 8, "8": 8,
    "sep": 9, "september": 9, "9": 9,
    "oct": 10, "october": 10, "10": 10,
    "nov": 11, "november": 11, "11": 11,
    "dec": 12, "december": 12, "12": 12,
}

DAYS_IN_MONTH = [0, 31, 29, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]


def _month_num(m: str) -> int:
    return MONTH_MAP.get(m.strip().lower(), 1)


def _month_day(m: int, d: int, y: int) -> int:
    if d < 1:
        return 1
    max_d = DAYS_IN_MONTH[m]
    if m == 2 and d == 29 and not (y % 4 == 0 and (y % 100 != 0 or y % 400 == 0)):
        max_d = 28
    return min(d, max_d)


def parse_pub_date(article_elem: ET.Element) -> Optional[date]:
    """Parse PubDate element from PubMed XML → date object with month/day or None."""
    pub_date_elem = article_elem.find(".//PubDate")
    if pub_date_elem is None:
        return None

    # Normal path: Year + Month + Day elements
    y_elem = pub_date_elem.find("Year")
    m_elem = pub_date_elem.find("Month")
    d_elem = pub_date_elem.find("Day")
    if y_elem is not None and y_elem.text:
        y = int(y_elem.text.strip())
        has_month = m_elem is not None and m_elem.text
        has_day = d_elem is not None and d_elem.text
        if not has_month and not has_day:
            return None  # Year-only, no better than what we already have
        m = _month_num(m_elem.text.strip()) if has_month else 1
        d = int(d_elem.text.strip()) if has_day else 1
        d = _month_day(m, d, y)
        return date(y, m, d)

    # MedlineDate format: "2024 Jan-Dec" → still only year precision
    medline = pub_date_elem.find("MedlineDate")
    if medline is not None and medline.text:
        import re
        m = re.match(r"(\d{4})", medline.text.strip())
        if m:
            return None  # year-level only

    return None


def parse_pmid(article_elem: ET.Element) -> Optional[int]:
    pmid_elem = article_elem.find("MedlineCitation/PMID")
    if pmid_elem is not None and pmid_elem.text:
        return int(pmid_elem.text.strip())
    return None


# ── main script ──

async def main():
    import argparse
    parser = argparse.ArgumentParser(description="Fix pub_date via NCBI efetch")
    parser.add_argument("--limit", type=int, default=0, help="Max records to fix (0=all)")
    parser.add_argument("--batch", type=int, default=200, help="Batch size per efetch req")
    parser.add_argument("--rate", type=float, default=3.0, help="Max req/sec")
    parser.add_argument("--dry-run", action="store_true", help="Don't update DB, just print")
    args = parser.parse_args()

    engine = create_async_engine(DATABASE_URL, pool_size=4)
    async_session = sessionmaker(engine, class_=AsyncSession)

    # 1. Get PMIDs with YYYY-01-01 dates
    async with async_session() as s:
        if args.dry_run:
            row = (await s.execute(
                text("SELECT COUNT(*) FROM global_literature WHERE pub_date IS NOT NULL AND EXTRACT(MONTH FROM pub_date)=1 AND EXTRACT(DAY FROM pub_date)=1")
            )).scalar()
            print(f"Records with YYYY-01-01: {row}")
            return

        rows = (await s.execute(
            text("""
                SELECT id, pmid FROM global_literature
                WHERE pub_date IS NOT NULL
                  AND EXTRACT(MONTH FROM pub_date)=1
                  AND EXTRACT(DAY FROM pub_date)=1
                  AND pmid IS NOT NULL
                ORDER BY pub_year ASC, pmid ASC
            """)
        )).all()

    total = len(rows)
    if args.limit and args.limit < total:
        rows = rows[:args.limit]
        total = args.limit

    print(f"Found {len(rows)} records to fix (showing dates for {min(5, len(rows))}):")
    for r in rows[:5]:
        print(f"  id={r.id} pmid={r.pmid}")

    # 2. Batch query efetch
    BATCH = args.batch
    INTERVAL = BATCH / args.rate  # seconds between batches

    client = httpx.AsyncClient(timeout=60.0)
    efetch_url = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"

    updated_count = 0
    error_count = 0

    for i in range(0, len(rows), BATCH):
        batch = rows[i:i + BATCH]
        pmids = [str(r.pmid) for r in batch]

        if i > 0:
            await asyncio.sleep(INTERVAL)

        try:
            resp = await client.get(
                efetch_url,
                params={"db": "pubmed", "id": ",".join(pmids), "retmode": "xml"},
            )
            resp.raise_for_status()
        except Exception as e:
            print(f"  batch {i//BATCH + 1}: HTTP error: {e}")
            error_count += len(batch)
            continue

        # Parse XML
        try:
            root = ET.fromstring(resp.content)
        except Exception as e:
            print(f"  batch {i//BATCH + 1}: XML parse error: {e}")
            error_count += len(batch)
            continue

        # Build pmid→pub_date map from response
        date_map: Dict[int, date] = {}
        for article in root.findall("PubmedArticle"):
            pmid = parse_pmid(article)
            if pmid is None:
                continue
            pd = parse_pub_date(article)
            if pd:
                date_map[pmid] = pd

        # Update DB
        batch_updates = 0
        async with async_session() as s:
            for r in batch:
                pd = date_map.get(r.pmid)
                if pd is None:
                    continue
                # skip if efetch also returns Jan-01 (no better date available)
                if pd.month == 1 and pd.day == 1:
                    continue
                await s.execute(
                    text("UPDATE global_literature SET pub_date = :pd, updated_at = NOW() WHERE id = :id AND pmid = :pmid"),
                    {"pd": pd, "id": r.id, "pmid": r.pmid},
                )
                batch_updates += 1
            await s.commit()

        updated_count += batch_updates
        pct = (i + BATCH) / total * 100 if total else 100
        print(f"  batch {i//BATCH + 1}/{(total + BATCH - 1)//BATCH}: +{batch_updates} updated ({pct:.0f}%)", flush=True)

    await client.aclose()
    await engine.dispose()

    print(f"\nDone. Updated: {updated_count}, Errors: {error_count}")


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