"""Download Alzheimer's disease papers from PubMed Central.

This module wraps the ``pubmed-stream`` library (https://pypi.org/project/pubmed-stream/)
to provide Alzheimer's-focused defaults and a convenience function for building a
comprehensive AD corpus across a curated set of research topics.
"""

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import Any, List, Optional

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Curated Alzheimer's disease search terms
# ---------------------------------------------------------------------------

AD_SEARCH_TERMS: List[str] = [
    "Alzheimer's disease pathogenesis",
    "Alzheimer disease amyloid beta plaques",
    "Alzheimer disease tau neurofibrillary tangles",
    "Alzheimer disease neuroinflammation microglia",
    "Alzheimer disease gut microbiome dysbiosis",
    "Alzheimer disease biomarkers cerebrospinal fluid",
    "Alzheimer disease treatment clinical trial",
    "Alzheimer disease APOE genetics risk",
    "Alzheimer disease cognitive decline progression",
    "Alzheimer disease blood-brain barrier",
    "Alzheimer disease synaptic plasticity",
    "Alzheimer disease gut brain axis",
    "Alzheimer disease probiotics prebiotics",
    "Alzheimer disease oxidative stress",
    "Alzheimer disease epigenetics",
]


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------


def download_papers(
    keyword: str,
    max_results: int = 100,
    output_dir: str | Path = "./publications",
    fmt: str = "text",
    api_key: Optional[str] = None,
    email: Optional[str] = None,
    use_concurrent: bool = True,
    max_workers: int = 5,
    rate_limit: Optional[float] = None,
    verbose: bool = False,
) -> Any:
    """Download PubMed Central articles matching *keyword* via pubmed-stream.

    Parameters
    ----------
    keyword:
        PubMed search query (e.g. ``"Alzheimer's disease gut microbiome"``).
    max_results:
        Maximum number of articles to download.
    output_dir:
        Base directory for downloaded JSON files. Articles are placed in a
        subdirectory named after the slugified keyword.
    fmt:
        Output format: ``"text"`` (default), ``"xml"``, or ``"both"``.
        All formats are saved as ``.json`` files.
    api_key:
        NCBI API key. Falls back to the ``NCBI_API_KEY`` environment variable.
        Increases the rate limit from 3 to 10 requests/second.
    email:
        Contact e-mail for NCBI. Falls back to ``NCBI_EMAIL`` env var.
    use_concurrent:
        Whether to use a thread pool for parallel downloads (default: ``True``).
    max_workers:
        Number of download threads when *use_concurrent* is ``True``.
    rate_limit:
        Seconds between API requests. ``None`` means auto (0.1 s with key,
        0.34 s without).
    verbose:
        Enable ``DEBUG``-level logging.

    Returns
    -------
    ``DownloadStats`` object from *pubmed-stream* with attributes:
    ``.successful``, ``.failed``, ``.skipped``, ``.total_found``,
    ``.duration_seconds``, ``.success_rate``.

    Raises
    ------
    ImportError
        If ``pubmed-stream`` is not installed.
    """
    try:
        from pubmed_stream import search_and_download
    except ImportError as exc:
        raise ImportError(
            "pubmed-stream is required for downloading papers.\n"
            "Install it with:  pip install pubmed-stream\n"
            "or install this package with full extras:  pip install neuromind[all]"
        ) from exc

    if verbose:
        logging.basicConfig(level=logging.DEBUG)

    # Resolve credentials from environment variables as a fallback
    resolved_key = api_key or os.environ.get("NCBI_API_KEY")
    resolved_email = email or os.environ.get("NCBI_EMAIL")

    logger.info("Downloading ≤%d papers for query: %r", max_results, keyword)

    kwargs: dict = dict(
        keyword=keyword,
        max_results=max_results,
        fmt=fmt,
        output_dir=Path(output_dir),
        use_concurrent=use_concurrent,
        max_workers=max_workers,
    )
    if resolved_key is not None:
        kwargs["api_key"] = resolved_key
    if resolved_email is not None:
        kwargs["email"] = resolved_email
    if rate_limit is not None:
        kwargs["rate_limit"] = rate_limit

    stats = search_and_download(**kwargs)

    logger.info(
        "Download complete — %d succeeded, %d failed, %d skipped (%.1fs)",
        stats.successful,
        stats.failed,
        stats.skipped,
        stats.duration_seconds,
    )
    return stats


def download_ad_corpus(
    topics: Optional[List[str]] = None,
    max_results_per_topic: int = 100,
    output_dir: str | Path = "./publications",
    api_key: Optional[str] = None,
    email: Optional[str] = None,
    **kwargs: Any,
) -> List[Any]:
    """Download a comprehensive Alzheimer's disease corpus.

    Iterates over *topics* (or the built-in :data:`AD_SEARCH_TERMS`) and
    downloads up to *max_results_per_topic* articles for each, depositing
    them into subdirectories under *output_dir*.

    Parameters
    ----------
    topics:
        List of search queries. Defaults to :data:`AD_SEARCH_TERMS`
        (15 curated Alzheimer's topics).
    max_results_per_topic:
        Maximum articles per topic (default: 100 → up to 1 500 total).
    output_dir:
        Root download directory.
    api_key:
        NCBI API key (or set ``NCBI_API_KEY`` env var).
    email:
        Contact e-mail for NCBI (or set ``NCBI_EMAIL`` env var).
    **kwargs:
        Additional keyword arguments forwarded to :func:`download_papers`.

    Returns
    -------
    List of ``DownloadStats`` objects, one per topic.
    """
    if topics is None:
        topics = AD_SEARCH_TERMS

    all_stats: List[Any] = []
    for i, topic in enumerate(topics, 1):
        logger.info("[%d/%d] %s", i, len(topics), topic)
        stats = download_papers(
            keyword=topic,
            max_results=max_results_per_topic,
            output_dir=output_dir,
            api_key=api_key,
            email=email,
            **kwargs,
        )
        all_stats.append(stats)

    total_ok = sum(s.successful for s in all_stats)
    total_fail = sum(s.failed for s in all_stats)
    logger.info(
        "Corpus download complete — %d articles downloaded, %d failed across %d topics",
        total_ok,
        total_fail,
        len(topics),
    )
    return all_stats
