"""Phase B — full-text retrieval of stage-1 candidates from legal open-access sources only:
Semantic Scholar openAccessPdf, Unpaywall (best_oa_location) by DOI, arXiv. Records every attempt in
data/metadata/retrieval_log.csv; PDFs stored in data/raw/pdfs/<record_id>.pdf (not committed).
Usage: python scripts/extraction/fetch_pdfs.py [--limit N]"""
import argparse, csv, hashlib, json, re, time, urllib.parse, urllib.request
from datetime import date
from pathlib import Path
import pandas as pd

ROOT = Path(__file__).resolve().parents[2]; META = ROOT / "data" / "metadata"; PDFS = ROOT / "data" / "raw" / "pdfs"; PDFS.mkdir(parents=True, exist_ok=True)
MAILTO = "simon-pierre.boucher@uqo.ca"
UA = f"Mozilla/5.0 (compatible; wp21-uqo meta-analysis; mailto:{MAILTO})"


def fetch(url, timeout=60):
    req = urllib.request.Request(url, headers={"User-Agent": UA, "Accept": "application/pdf,*/*"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        data = r.read(); ctype = r.headers.get("Content-Type", "")
    return data, ctype


def unpaywall(doi):
    """All OA locations from Unpaywall, repositories first, pdf URLs before landing pages. Returns list of (label, url)."""
    try:
        with urllib.request.urlopen(urllib.request.Request(f"https://api.unpaywall.org/v2/{urllib.parse.quote(doi)}?email={MAILTO}", headers={"User-Agent": UA}), timeout=30) as r:
            d = json.load(r)
    except Exception:
        return []
    locs = d.get("oa_locations") or []
    locs = sorted(locs, key=lambda l: (0 if l.get("host_type") == "repository" else 1))
    out = []
    for l in locs:
        for key in ("url_for_pdf", "url"):
            u = l.get(key)
            if u and (f"unpaywall:{l.get('host_type')}", u) not in out:
                out.append((f"unpaywall:{l.get('host_type')}", u))
    return out[:6]


def pdf_link_from_html(html, base):
    m = re.search(r'<meta[^>]+name="citation_pdf_url"[^>]+content="([^"]+)"', html, re.I) or re.search(r'content="([^"]+)"[^>]+name="citation_pdf_url"', html, re.I)
    if m:
        return urllib.parse.urljoin(base, m.group(1))
    m = re.search(r'href="([^"]+\.pdf[^"]*)"', html, re.I)
    return urllib.parse.urljoin(base, m.group(1)) if m else None


def is_pdf(data):
    return data[:5] == b"%PDF-"


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("--limit", type=int, default=100000); ap.add_argument("--input", default=str(META / "stage1_candidates.csv")); a = ap.parse_args()
    cand = pd.read_csv(a.input)
    logf = META / "retrieval_log.csv"; done = set()
    if logf.exists():
        done = {r["record_id"] for r in csv.DictReader(open(logf)) if r["status"] in ("ok", "no_oa_url")}
    with open(logf, "a", newline="") as lf:
        log = csv.writer(lf)
        if lf.tell() == 0:
            log.writerow(["date", "record_id", "doi", "source_tried", "url", "status", "bytes", "sha256"])
        n_ok = 0
        for _, r in cand.head(a.limit).iterrows():
            rid = str(r.get("record_id") or r.sid)
            if rid in done or (PDFS / f"{rid}.pdf").exists():
                continue
            urls = []
            if isinstance(r.get("oa_pdf"), str) and r.oa_pdf.startswith("http"):
                urls.append(("s2_oa", r.oa_pdf))
            if isinstance(r.get("arxiv"), str) and r.arxiv:
                urls.append(("arxiv", f"https://arxiv.org/pdf/{r.arxiv}"))
            if isinstance(r.get("doi"), str) and r.doi:
                urls += unpaywall(r.doi)
            if not urls:
                log.writerow([str(date.today()), rid, r.get("doi", ""), "none", "", "no_oa_url", 0, ""]); continue
            got = False
            for src, u in urls:
                try:
                    data, ctype = fetch(u)
                    if not is_pdf(data) and b"<html" in data[:2000].lower():
                        link = pdf_link_from_html(data.decode("utf-8", "ignore"), u)
                        if link and link != u:
                            try:
                                data, ctype = fetch(link); u = link; src = src + "+html_pdf_link"
                            except Exception:
                                pass
                    if is_pdf(data) and len(data) > 10_000:
                        (PDFS / f"{rid}.pdf").write_bytes(data)
                        log.writerow([str(date.today()), rid, r.get("doi", ""), src, u, "ok", len(data), hashlib.sha256(data).hexdigest()[:16]]); got = True; n_ok += 1; break
                    log.writerow([str(date.today()), rid, r.get("doi", ""), src, u, f"not_pdf:{ctype[:30]}", len(data), ""])
                except Exception as e:
                    log.writerow([str(date.today()), rid, r.get("doi", ""), src, u, f"error:{type(e).__name__}", 0, ""])
                time.sleep(0.5)
            lf.flush()
        print("retrieved", n_ok)


if __name__ == "__main__":
    main()