"""
rag_data_loader.py

Study-aware PubMed ingestion utility for the RAG Gene Discovery pipeline.

This module fetches PubMed abstracts, metadata, and (when available) open-access
PDFs for a given search query and organizes them into a consistent, study-scoped
directory structure. The output of this loader serves as the canonical input for
downstream embedding, indexing, and retrieval components.

Directory layout (resource-first, study-second):

    PubMed/
      Abstracts/<study>/
      Metadata/<study>/
      PDFs/<study>/
      Index/<study>/

This layout is intentionally aligned with `embedding_engine.py`, which expects
abstracts under `Abstracts/<study>` and writes FAISS indices under `Index/<study>`.

Key features:
- PubMed search via NCBI Entrez (with pagination and retry support)
- Study-level isolation to support multiple case studies in parallel
- Incremental ingestion with resume support (skips already processed PMIDs)
- Optional best-effort download of open-access PDFs from PubMed Central (PMC)
- Clear provenance and reproducibility through study-scoped folders

This module focuses exclusively on **data ingestion and organization**.
Vector embedding generation and FAISS index construction are intentionally handled
by separate modules.

Command-line usage:

    python -m ragbio.utils.rag_data_loader \
        --study Alzheimer_CaseStudy \
        --search "Alzheimer Disease AND therapy" \
        --retmax 500 \
        --retstart 0

Required arguments:
- --study   : Logical study or dataset name (used as a folder namespace)
- --search  : PubMed search query

Optional arguments:
- --retmax    : Number of records per batch (default: 500)
- --retstart  : Starting offset for PubMed pagination (default: 0)
- --query     : Alias for --search (kept for backward compatibility)

Author:
    Manish Kumar
"""

from __future__ import annotations

import argparse
import json
import os
import time
from pathlib import Path
from typing import Dict, List, Set

import requests
from bs4 import BeautifulSoup
from Bio import Entrez
from urllib.error import HTTPError, URLError


# ----------------------------
# Paths (study-scoped)
# ----------------------------
def resolve_workspace_dir() -> Path:
    # Docker convention
    if Path("/workspace").exists():
        return Path("/workspace")
    return Path.cwd()


def resolve_pubmed_base_path(workspace_dir: Path) -> Path:
    # Prefer env var if set (your config.py uses this)
    env_base = os.getenv("PUBMED_BASE_PATH")
    if env_base:
        return Path(env_base)
    # Default similar to your config.py, but safe here too
    return workspace_dir / "data" / "PubMed"


def get_study_dirs(base_path: Path, study: str) -> Dict[str, Path]:
    """
    Study-aware folders that are CONSISTENT with embedding_engine.py

    Layout:
      PubMed/
        Abstracts/<study>/
        Metadata/<study>/
        PDFs/<study>/
        Index/<study>/
    """
    abstracts = base_path / "Abstracts" / study
    metadata  = base_path / "Metadata"  / study
    pdfs      = base_path / "PDFs"      / study
    index     = base_path / "Index"     / study

    for p in (abstracts, metadata, pdfs, index):
        p.mkdir(parents=True, exist_ok=True)

    return {
        "abstracts": abstracts,
        "metadata": metadata,
        "pdfs": pdfs,
        "index": index,
        "index_file": index / "pubmed_index.faiss",
        "id_map_file": index / "pmid_map.json",
    }

# ----------------------------
# PubMed helpers
# ----------------------------
def get_total_records(term: str) -> int:
    handle = Entrez.esearch(db="pubmed", term=term, retmax=1)
    record = Entrez.read(handle)
    return int(record.get("Count", 0))


def fetch_pubmed_ids(term: str, retmax: int, retstart: int, retries: int = 3, delay: int = 5) -> List[str]:
    for attempt in range(1, retries + 1):
        try:
            handle = Entrez.esearch(db="pubmed", term=term, retmax=retmax, retstart=retstart)
            record = Entrez.read(handle)
            return record.get("IdList", [])
        except (HTTPError, URLError, RuntimeError) as e:
            print(f"[WARN] Entrez.esearch attempt {attempt} failed: {e}")
            time.sleep(delay)
    print("[ERROR] Failed to fetch PubMed IDs after multiple attempts.")
    return []


def fetch_abstract_and_metadata(pmid: str) -> Dict:
    handle = Entrez.efetch(db="pubmed", id=pmid, retmode="xml")
    records = Entrez.read(handle)

    article = records["PubmedArticle"][0]["MedlineCitation"]["Article"]
    title = article.get("ArticleTitle", "")

    abstract = ""
    if "Abstract" in article and "AbstractText" in article["Abstract"]:
        # AbstractText sometimes contains strings and sometimes objects; coerce to str safely
        abstract = " ".join(str(x) for x in article["Abstract"].get("AbstractText", []))

    authors: List[str] = []
    if "AuthorList" in article:
        for a in article["AuthorList"]:
            name = f"{a.get('LastName', '')} {a.get('ForeName', '')}".strip()
            if name:
                authors.append(name)

    return {"pmid": pmid, "title": str(title), "abstract": abstract, "authors": authors}


def save_json(data: Dict, out_path: Path) -> Path:
    out_path.parent.mkdir(parents=True, exist_ok=True)
    with out_path.open("w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False, indent=2)
    return out_path


def try_download_pdf(pmid: str, pdf_folder: Path, timeout: int = 10) -> bool:
    """
    Best-effort: searches PMC for a PDF link.
    Not all PubMed entries have OA PDFs.
    """
    try:
        pmc_url = f"https://www.ncbi.nlm.nih.gov/pmc/?term={pmid}"
        response = requests.get(pmc_url, timeout=timeout)
        response.raise_for_status()

        soup = BeautifulSoup(response.text, "html.parser")
        pdf_link = soup.find("a", string="PDF")
        if not pdf_link:
            return False

        pdf_href = pdf_link.get("href")
        if not pdf_href:
            return False
        if pdf_href.startswith("/"):
            pdf_href = "https://www.ncbi.nlm.nih.gov" + pdf_href

        pdf_response = requests.get(pdf_href, timeout=timeout)
        pdf_response.raise_for_status()

        pdf_path = pdf_folder / f"{pmid}.pdf"
        with pdf_path.open("wb") as f:
            f.write(pdf_response.content)

        print(f"[PDF] Downloaded for PMID {pmid}")
        return True

    except Exception as e:
        print(f"[WARN] PDF not available for PMID {pmid}: {e}")
        return False


def get_processed_pmids(abstract_folder: Path) -> Set[str]:
    processed: Set[str] = set()
    if not abstract_folder.exists():
        return processed
    for fp in abstract_folder.glob("*.json"):
        processed.add(fp.stem)
    return processed


# ----------------------------
# CLI
# ----------------------------
def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        prog="rag_data_loader",
        description="Fetch PubMed abstracts/metadata/PDFs into study-scoped folders.",
    )

    # Keep your original flags working:
    p.add_argument("--search", dest="search", required=True, help="PubMed search term (e.g. \"Alzheimer Disease AND therapy\")")
    p.add_argument("--retmax", dest="retmax", type=int, default=500, help="Records per batch (default: 500)")
    p.add_argument("--retstart", dest="retstart", type=int, default=0, help="Start offset for pagination (default: 0)")
    p.add_argument("--study", dest="study", required=True, help="Study folder name (e.g. Alzheimer_CaseStudy)")

    # Also accept the older/newer naming you accidentally introduced:
    # (If someone calls it with --query, treat it as --search)
    p.add_argument("--query", dest="query", default=None, help="Alias for --search (kept for compatibility)")

    return p.parse_args()


def main() -> int:
    args = parse_args()

    search_term = args.search
    if args.query and not args.search:
        search_term = args.query
    if args.query and args.search and args.query != args.search:
        print("[WARN] Both --query and --search provided; using --search.")

    # Entrez config
    Entrez.api_base = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/"
    Entrez.email = os.getenv("NCBI_EMAIL", "mandecent.gupta@gmail.com")

    workspace = resolve_workspace_dir()
    base_path = resolve_pubmed_base_path(workspace)
    dirs = get_study_dirs(base_path, args.study)

    print(f"Workspace directory: {workspace}")
    print(f"Base path: {base_path}")
    print(f"Abstract folder: {dirs['abstracts']}")
    print(f"Metadata folder: {dirs['metadata']}")
    print(f"PDF folder: {dirs['pdfs']}")
    print(f"Index folder: {dirs['index']}")
    print(f"Index file: {dirs['index_file']}")
    print(f"ID map file: {dirs['id_map_file']}")

    total = get_total_records(search_term)
    print(f"\n[INFO] Total available records for '{search_term}': {total}")

    processed = get_processed_pmids(dirs["abstracts"])
    print(f"[INFO] Already processed in this study: {len(processed)} PMIDs.")

    retstart = args.retstart
    retmax = args.retmax

    while retstart < total:
        pmids = fetch_pubmed_ids(search_term, retmax=retmax, retstart=retstart)
        if not pmids:
            break

        print(f"\n[INFO] Fetching batch starting at {retstart}, size={len(pmids)}")

        for pmid in pmids:
            if pmid in processed:
                # keep logs short
                continue

            try:
                record = fetch_abstract_and_metadata(pmid)

                # abstract json (light)
                save_json(
                    {"pmid": pmid, "title": record.get("title", ""), "abstract": record.get("abstract", "")},
                    dirs["abstracts"] / f"{pmid}.json",
                )

                # metadata json (richer)
                save_json(record, dirs["metadata"] / f"{pmid}_meta.json")

                # optional pdf
                try_download_pdf(pmid, dirs["pdfs"])

                processed.add(pmid)
                title_preview = (record.get("title") or "")[:80]
                print(f"[OK] PMID {pmid} - {title_preview}")

            except Exception as e:
                print(f"[ERROR] Problem with PMID {pmid}: {e}")

        retstart += retmax
        print(f"[INFO] Completed batch. Next start index: {retstart}")
        time.sleep(2)

    print(f"\n[INFO] Done. Study '{args.study}' processed PMIDs: {len(processed)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
