"""
scripts/download_corpus.py
--------------------------
Download 100–200 scientific papers from arXiv in the NLP/LLM/RAG domain.
Saves PDFs to data/papers/ and metadata to data/metadata.jsonl.

Usage:
    python scripts/download_corpus.py --max_results 150 --output_dir data/papers
"""

import argparse
import json
import time
import urllib.request
import urllib.parse
from pathlib import Path
import xml.etree.ElementTree as ET


ARXIV_QUERIES = [
    "retrieval augmented generation large language models",
    "RAG question answering transformer",
    "dense retrieval natural language processing",
    "LLM hallucination mitigation grounding",
    "FAISS vector search semantic similarity",
    "sentence transformers embeddings NLP",
    "scientific question answering benchmark",
    "knowledge-intensive NLP tasks",
]

ARXIV_API = "https://export.arxiv.org/api/query"
NS = {"atom": "http://www.w3.org/2005/Atom", "arxiv": "http://arxiv.org/schemas/atom"}


def search_arxiv(query: str, max_results: int = 25, start: int = 0) -> list[dict]:
    params = urllib.parse.urlencode({
        "search_query": f"all:{query}",
        "start": start,
        "max_results": max_results,
        "sortBy": "relevance",
        "sortOrder": "descending",
    })
    url = f"{ARXIV_API}?{params}"
    with urllib.request.urlopen(url, timeout=30) as resp:
        xml_data = resp.read()
    root = ET.fromstring(xml_data)
    papers = []
    for entry in root.findall("atom:entry", NS):
        arxiv_id_raw = entry.findtext("atom:id", default="", namespaces=NS)
        arxiv_id = arxiv_id_raw.split("/abs/")[-1].strip()
        title = (entry.findtext("atom:title", default="", namespaces=NS) or "").strip().replace("\n", " ")
        abstract = (entry.findtext("atom:summary", default="", namespaces=NS) or "").strip().replace("\n", " ")
        authors = [
            a.findtext("atom:name", default="", namespaces=NS)
            for a in entry.findall("atom:author", NS)
        ]
        published = entry.findtext("atom:published", default="", namespaces=NS)
        year = published[:4] if published else "unknown"
        # PDF link
        pdf_url = None
        for link in entry.findall("atom:link", NS):
            if link.get("type") == "application/pdf":
                pdf_url = link.get("href")
                break
        if pdf_url is None and arxiv_id:
            pdf_url = f"https://arxiv.org/pdf/{arxiv_id}"
        papers.append({
            "arxiv_id": arxiv_id,
            "title": title,
            "authors": authors,
            "year": year,
            "abstract": abstract,
            "pdf_url": pdf_url,
        })
    return papers


def download_pdf(pdf_url: str, dest_path: Path) -> bool:
    if dest_path.exists():
        return True  # already downloaded
    try:
        headers = {"User-Agent": "RAG-ScientificCorpus/1.0 (research project)"}
        req = urllib.request.Request(pdf_url, headers=headers)
        with urllib.request.urlopen(req, timeout=60) as resp:
            dest_path.write_bytes(resp.read())
        return True
    except Exception as e:
        print(f"  ✗ Failed {pdf_url}: {e}")
        return False


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--max_results", type=int, default=150)
    parser.add_argument("--output_dir", type=str, default="data/papers")
    parser.add_argument("--metadata_file", type=str, default="data/metadata.jsonl")
    parser.add_argument("--sleep", type=float, default=3.0, help="Seconds between API calls")
    args = parser.parse_args()

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    meta_path = Path(args.metadata_file)

    seen_ids: set[str] = set()
    all_papers: list[dict] = []

    per_query = max(10, args.max_results // len(ARXIV_QUERIES))

    for query in ARXIV_QUERIES:
        if len(all_papers) >= args.max_results:
            break
        print(f"\n🔍 Querying: '{query}'")
        try:
            results = search_arxiv(query, max_results=per_query)
        except Exception as e:
            print(f"  API error: {e}")
            time.sleep(args.sleep * 2)
            continue

        for paper in results:
            if paper["arxiv_id"] in seen_ids or not paper["arxiv_id"]:
                continue
            seen_ids.add(paper["arxiv_id"])
            all_papers.append(paper)
            if len(all_papers) >= args.max_results:
                break

        print(f"  → {len(all_papers)} unique papers so far")
        time.sleep(args.sleep)

    print(f"\n📥 Downloading {len(all_papers)} PDFs to {out_dir}/ ...")
    downloaded = 0
    for i, paper in enumerate(all_papers):
        safe_id = paper["arxiv_id"].replace("/", "_").replace(".", "_")
        dest = out_dir / f"{safe_id}.pdf"
        print(f"  [{i+1}/{len(all_papers)}] {paper['title'][:60]}...")
        ok = download_pdf(paper["pdf_url"], dest)
        if ok:
            paper["local_path"] = str(dest)
            downloaded += 1
        else:
            paper["local_path"] = None
        time.sleep(1.0)  # be polite to arXiv

    # Write metadata
    with open(meta_path, "w", encoding="utf-8") as f:
        for p in all_papers:
            f.write(json.dumps(p, ensure_ascii=False) + "\n")

    print(f"\n✅ Done. {downloaded}/{len(all_papers)} PDFs downloaded.")
    print(f"   Metadata saved to: {meta_path}")


if __name__ == "__main__":
    main()