src/literature.py

"""
PubMed E-utilities client. No API key needed (3 req/sec).
Falls back gracefully when rate-limited.
"""

import time
import requests
import xml.etree.ElementTree as ET
from typing import Optional
from src.cache import get_cached, set_cached


EUTILS_BASE = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils"
_last_req = 0.0


def _rate_limit():
    global _last_req
    wait = 0.35 - (time.time() - _last_req)
    if wait > 0:
        time.sleep(wait)
    _last_req = time.time()


def _search(query: str, max_results: int = 3) -> Optional[list[str]]:
    _rate_limit()
    try:
        resp = requests.get(f"{EUTILS_BASE}/esearch.fcgi", params={
            "db": "pubmed", "retmode": "json", "retmax": max_results,
            "sort": "relevance", "term": query,
        }, timeout=10)
        if resp.status_code == 429:
            return None
        resp.raise_for_status()
        return resp.json().get("esearchresult", {}).get("idlist", [])
    except Exception as e:
        print(f"PubMed search error: {e}")
        return None


def _fetch(pmids: list[str]) -> list[dict]:
    if not pmids:
        return []
    _rate_limit()
    try:
        resp = requests.get(f"{EUTILS_BASE}/efetch.fcgi", params={
            "db": "pubmed", "id": ",".join(pmids),
            "retmode": "xml", "rettype": "abstract",
        }, timeout=10)
        resp.raise_for_status()

        articles = []
        root = ET.fromstring(resp.text)
        for elem in root.findall(".//PubmedArticle"):
            a = _parse(elem)
            if a:
                articles.append(a)
        return articles
    except Exception as e:
        print(f"PubMed fetch error: {e}")
        return []


def _parse(elem) -> Optional[dict]:
    try:
        med = elem.find(".//MedlineCitation")
        art = med.find(".//Article")
        pmid = med.findtext("PMID", "")
        title = art.findtext("ArticleTitle", "")

        authors = []
        al = art.find("AuthorList")
        if al is not None:
            for a in al.findall("Author")[:3]:
                last = a.findtext("LastName", "")
                init = a.findtext("Initials", "")
                if last:
                    authors.append(f"{last} {init}".strip())

        author_str = ", ".join(authors)
        if len(al.findall("Author") if al is not None else []) > 3:
            author_str += " et al."

        j = art.find("Journal")
        journal = j.findtext("ISOAbbreviation", "") if j is not None else ""
        pd = j.find(".//PubDate") if j is not None else None
        year = pd.findtext("Year", "") if pd is not None else ""

        abstract_parts = []
        ab = art.find("Abstract")
        if ab is not None:
            for t in ab.findall("AbstractText"):
                label = t.get("Label", "")
                content = (t.text or "")[:400]
                abstract_parts.append(f"{label}: {content}" if label else content)

        doi = ""
        for id_el in elem.findall(".//ArticleId"):
            if id_el.get("IdType") == "doi":
                doi = id_el.text or ""
                break

        return {
            "pmid": pmid, "title": title, "authors": author_str,
            "journal": journal, "year": year,
            "abstract": " ".join(abstract_parts)[:800],
            "doi": doi, "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/",
        }
    except Exception:
        return None


def search_food_health(ingredient: str, health_aspect: str = "",
                       max_papers: int = 2) -> list[dict]:
    """Search PubMed for studies on a food's health effects. Cached."""
    if health_aspect:
        query = f'({ingredient}) AND ({health_aspect}) AND (health OR nutrition)'
    else:
        query = f'({ingredient}) AND (health effects OR nutritional)'
    query += ' AND (Review[pt] OR Meta-Analysis[pt])'

    cache_key = f"lit:{ingredient}:{health_aspect}"
    cached = get_cached("literature_cache", cache_key)
    if cached is not None:
        return cached

    pmids = _search(query, max_results=max_papers)
    if not pmids:
        # broader retry without review filter
        pmids = _search(query.replace(' AND (Review[pt] OR Meta-Analysis[pt])', ''),
                        max_results=max_papers)
    if not pmids:
        return []

    articles = _fetch(pmids)
    set_cached("literature_cache", cache_key, articles, ttl_days=7)
    return articles


def format_citation(a: dict) -> str:
    parts = []
    if a.get("authors"): parts.append(a["authors"])
    if a.get("year"): parts.append(f"({a['year']})")
    if a.get("title"): parts.append(f"{a['title'].rstrip('.')}.")
    if a.get("journal"): parts.append(f"{a['journal']}.")
    if a.get("doi"): parts.append(f"DOI: {a['doi']}")
    return " ".join(parts)


def lookup_literature(ingredients: list[str], health_goal: str = "",
                      papers_per: int = 2) -> tuple[dict, int]:
    """Look up literature for ingredients. Returns (results, failure_count)."""
    results = {}
    failures = 0
    for ing in ingredients:
        articles = search_food_health(ing.strip(), health_goal, papers_per)
        if articles:
            results[ing] = articles
        else:
            failures += 1
            results[ing] = []
    return results, failures