#!/usr/bin/env python3
"""Scan PubMed baseline XML files and verify oncology coverage in DB.

Usage:
    cd backend
    python scripts/scan_baseline_coverage.py /data/2026baseline
    python scripts/scan_baseline_coverage.py /data/2026baseline --max-files 10
    python scripts/scan_baseline_coverage.py /data/2026baseline --csv coverage.csv
    python scripts/scan_baseline_coverage.py /data/2026baseline --detail
"""

import argparse
import asyncio
import csv
import gzip
import sys
import time
from pathlib import Path

_SCRIPT_DIR = str(Path(__file__).parent)
_PROJECT_ROOT = str(Path(__file__).parent.parent)
if _PROJECT_ROOT not in sys.path:
    sys.path.insert(0, _PROJECT_ROOT)
if _SCRIPT_DIR not in sys.path:
    sys.path.insert(0, _SCRIPT_DIR)

from lxml import etree
from sqlalchemy import select

from app.db import async_session
from app.models.literature import GlobalLiterature
from app.services.mesh_filter import get_mesh_filter
from pubmed_baseline import _extract_article, _is_oncology


async def scan_file(xml_path: Path, batch_size: int = 500) -> dict:
    """Scan single .xml.gz, collecting oncology PMIDs and checking DB coverage."""
    stats = {
        'file': xml_path.name,
        'total': 0,
        'oncology': 0,
        'covered': 0,
        'missing': 0,
        'missing_pmids': [],
        'error': None,
    }
    oncology_pmids: set[int] = set()

    # Phase 1: Parse XML, collect oncology PMIDs
    try:
        with gzip.open(xml_path, 'rb') as f:
            context = etree.iterparse(f, events=('end',), tag='PubmedArticle')
            for _event, elem in context:
                article = _extract_article(elem)
                elem.clear()
                while elem.getprevious() is not None:
                    del elem.getparent()[0]
                if article is None:
                    continue
                stats['total'] += 1
                if _is_oncology(article):
                    stats['oncology'] += 1
                    oncology_pmids.add(article['pmid'])
    except Exception as e:
        stats['error'] = f"XML parse failed: {e}"
        return stats

    if not oncology_pmids:
        return stats

    # Phase 2: Batch-query DB for coverage
    async with async_session() as db:
        pmid_list = sorted(oncology_pmids)
        for i in range(0, len(pmid_list), batch_size):
            chunk = pmid_list[i:i + batch_size]
            result = await db.execute(
                select(GlobalLiterature.pmid).where(
                    GlobalLiterature.pmid.in_(chunk),
                    GlobalLiterature.source == 'pubmed_ftp'
                )
            )
            existing = {row[0] for row in result.all()}
            stats['covered'] += len(existing)
            missing_in_chunk = set(chunk) - existing
            stats['missing_pmids'].extend(sorted(missing_in_chunk))

    stats['missing'] = len(stats['missing_pmids'])
    return stats


async def main():
    parser = argparse.ArgumentParser(
        description='Scan baseline XML files and verify oncology PMIDs in DB'
    )
    parser.add_argument('dir', help='Directory containing .xml.gz baseline files')
    parser.add_argument('--max-files', type=int, default=0,
                        help='Max files to scan (0=all)')
    parser.add_argument('--csv', default='',
                        help='Output CSV path (e.g. coverage.csv)')
    parser.add_argument('--batch-size', type=int, default=500,
                        help='DB query batch size')
    parser.add_argument('--detail', action='store_true',
                        help='List missing PMIDs in summary')
    args = parser.parse_args()

    xml_dir = Path(args.dir)
    if not xml_dir.is_dir():
        print(f"ERROR: directory not found: {args.dir}", file=sys.stderr)
        sys.exit(1)

    files = sorted(xml_dir.glob('*.xml.gz'))
    if not files:
        print(f"No .xml.gz files found in {args.dir}")
        sys.exit(0)

    if args.max_files > 0:
        files = files[:args.max_files]

    print(f"Found {len(files)} XML files to scan\n"
          f"{'File':<34} {'Total':>8} {'Oncology':>10} {'Covered':>8} "
          f"{'Missing':>8} {'%':>6}")
    print('-' * 75)

    all_results: list[dict] = []
    all_missing: list[int] = []
    total_total = total_oncology = total_covered = 0
    start = time.time()

    # Warm up mesh_filter singleton (loads mtrees into cache once)
    get_mesh_filter()

    for f in files:
        result = await scan_file(f, args.batch_size)
        all_results.append(result)

        pct = (result['covered'] / result['oncology'] * 100) if result['oncology'] > 0 else 0
        suffix = f"  [ERR: {result['error']}]" if result['error'] else ""
        print(f"{result['file']:<34} {result['total']:>8,} {result['oncology']:>10,} "
              f"{result['covered']:>8,} {result['missing']:>8,} {pct:>5.1f}%{suffix}")

        total_total += result['total']
        total_oncology += result['oncology']
        total_covered += result['covered']
        all_missing.extend(result['missing_pmids'])

    elapsed = time.time() - start
    overall_pct = (total_covered / total_oncology * 100) if total_oncology > 0 else 0

    # Summary
    print('=' * 75)
    print(f"{'TOTAL':<34} {total_total:>8,} {total_oncology:>10,} "
          f"{total_covered:>8,} {len(all_missing):>8,} {overall_pct:>5.1f}%")
    print(f"\nFiles scanned: {len(files)}  |  Time: {elapsed:.1f}s  |  "
          f"Total articles: {total_total:,}  |  Oncology: {total_oncology:,}  |  "
          f"Missing from DB: {len(all_missing):,}")

    if total_oncology == 0:
        print("\nNo oncology articles found — check mtrees file and keywords.")
    elif overall_pct == 100.0:
        print("\nALL oncology articles are covered in the database.")
    else:
        print(f"\nCoverage: {overall_pct:.1f}% — {len(all_missing):,} oncology PMIDs missing.")

    if args.detail and all_missing:
        limit = min(50, len(all_missing))
        print(f"\nFirst {limit} missing PMIDs:")
        for pmid in all_missing[:limit]:
            print(f"  {pmid}")
        if len(all_missing) > limit:
            print(f"  ... and {len(all_missing) - limit} more")

    # CSV output
    if args.csv:
        with open(args.csv, 'w', newline='') as f:
            w = csv.writer(f)
            w.writerow(['file', 'total_articles', 'oncology_articles',
                       'covered', 'missing', 'coverage_pct', 'error'])
            for r in all_results:
                pct_r = (r['covered'] / r['oncology'] * 100) if r['oncology'] > 0 else 0
                w.writerow([
                    r['file'], r['total'], r['oncology'],
                    r['covered'], r['missing'], round(pct_r, 1), r['error'] or ''
                ])
        print(f"\nCSV written to: {args.csv}")

    if all_missing:
        sys.exit(1)


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