pubmed_watcher.py

#!/usr/bin/env python3
"""
PubMed Watcher — Automated alerting for new PubMed papers.

Watches configured queries, tracks seen papers, and reports only NEW results.
Pure Python 3 stdlib — no pip dependencies required.

Usage:
    python pubmed_watcher.py watch [--days N] [--json] [--markdown]
    python pubmed_watcher.py add "epigenetic reprogramming aging"
    python pubmed_watcher.py list
    python pubmed_watcher.py remove <id>
    python pubmed_watcher.py history [--limit N]
"""

import argparse
import json
import os
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
import xml.etree.ElementTree as ET
from datetime import datetime, timedelta, timezone
from pathlib import Path

# ---------------------------------------------------------------------------
# Paths
# ---------------------------------------------------------------------------
SCRIPT_DIR = Path(__file__).resolve().parent
CONFIG_PATH = SCRIPT_DIR / "config.json"
SEEN_PATH = SCRIPT_DIR / "seen.json"
HISTORY_PATH = SCRIPT_DIR / "history.json"

# ---------------------------------------------------------------------------
# NCBI E-utilities
# ---------------------------------------------------------------------------
ESEARCH_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"
EFETCH_URL = "https://eutils.ncbi.nlm.nih.gov/entrez/eutils/efetch.fcgi"
API_DELAY = 0.5  # seconds between requests (NCBI rate-limit friendly)

# ---------------------------------------------------------------------------
# Default queries
# ---------------------------------------------------------------------------
DEFAULT_QUERIES = [
    {"id": 1, "query": "epigenetic reprogramming aging longevity", "added": None},
    {"id": 2, "query": "immunosenescence aging immune system", "added": None},
    {"id": 3, "query": "senolytics senescent cells therapy", "added": None},
    {"id": 4, "query": "AI drug discovery clinical trial", "added": None},
    {"id": 5, "query": "biological age clock methylation", "added": None},
]


# ===================================================================
# State helpers
# ===================================================================

def _load_json(path: Path, default=None):
    """Load a JSON file, returning *default* if missing or corrupt."""
    if not path.exists():
        return default if default is not None else {}
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except (json.JSONDecodeError, OSError):
        return default if default is not None else {}


def _save_json(path: Path, data):
    """Atomically write JSON (write-then-rename)."""
    tmp = path.with_suffix(".tmp")
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
        f.write("\n")
    tmp.replace(path)


def load_config() -> dict:
    """Return config dict with at least {queries: [...], next_id: int}."""
    cfg = _load_json(CONFIG_PATH, default=None)
    if cfg is None or "queries" not in cfg:
        now = datetime.now(timezone.utc).isoformat() + "Z"
        queries = []
        for q in DEFAULT_QUERIES:
            queries.append({**q, "added": now})
        cfg = {"queries": queries, "next_id": len(queries) + 1}
        _save_json(CONFIG_PATH, cfg)
    return cfg


def save_config(cfg: dict):
    _save_json(CONFIG_PATH, cfg)


def load_seen() -> dict:
    """Return {pmid_str: first_seen_iso, ...}."""
    return _load_json(SEEN_PATH, default={})


def save_seen(seen: dict):
    _save_json(SEEN_PATH, seen)


def load_history() -> list:
    return _load_json(HISTORY_PATH, default=[])


def save_history(history: list):
    _save_json(HISTORY_PATH, history)


# ===================================================================
# PubMed API helpers
# ===================================================================

def _api_get(url: str, params: dict, retries: int = 3) -> bytes:
    """GET with retries and polite delay."""
    qs = urllib.parse.urlencode(params)
    full = f"{url}?{qs}"
    last_err = None
    for attempt in range(retries):
        try:
            req = urllib.request.Request(full, headers={"User-Agent": "pubmed-watcher/1.0"})
            with urllib.request.urlopen(req, timeout=30) as resp:
                return resp.read()
        except (urllib.error.URLError, urllib.error.HTTPError, OSError) as exc:
            last_err = exc
            if attempt < retries - 1:
                time.sleep(1 + attempt)
    raise ConnectionError(f"Failed after {retries} attempts: {last_err}")


def search_pmids(query: str, days: int = 7, retmax: int = 50) -> list[str]:
    """ESearch: return list of PMID strings for *query* within *days*."""
    mindate = (datetime.now(timezone.utc) - timedelta(days=days)).strftime("%Y/%m/%d")
    maxdate = datetime.now(timezone.utc).strftime("%Y/%m/%d")
    params = {
        "db": "pubmed",
        "term": query,
        "retmax": retmax,
        "usehistory": "n",
        "datetype": "edat",  # Entrez date (when added to PubMed)
        "mindate": mindate,
        "maxdate": maxdate,
        "retmode": "xml",
    }
    data = _api_get(ESEARCH_URL, params)
    root = ET.fromstring(data)
    id_list = root.find("IdList")
    if id_list is None:
        return []
    return [el.text for el in id_list.findall("Id") if el.text]


def fetch_details(pmids: list[str]) -> list[dict]:
    """EFetch: return list of paper dicts for given PMIDs."""
    if not pmids:
        return []
    params = {
        "db": "pubmed",
        "id": ",".join(pmids),
        "retmode": "xml",
        "rettype": "abstract",
    }
    data = _api_get(EFETCH_URL, params)
    root = ET.fromstring(data)
    papers = []
    for article_node in root.findall(".//PubmedArticle"):
        papers.append(_parse_article(article_node))
    return papers


def _parse_article(node) -> dict:
    """Extract key metadata from a PubmedArticle XML element."""

    def _text(el, path, default=""):
        found = el.find(path)
        return found.text.strip() if found is not None and found.text else default

    medline = node.find("MedlineCitation")
    article = medline.find("Article") if medline is not None else None

    pmid = _text(medline, "PMID") if medline is not None else ""
    title = ""
    if article is not None:
        title_el = article.find("ArticleTitle")
        if title_el is not None:
            # ArticleTitle can contain mixed content (italic etc.)
            title = "".join(title_el.itertext()).strip()

    # Authors
    authors = []
    author_list = article.find("AuthorList") if article is not None else None
    if author_list is not None:
        for au in author_list.findall("Author"):
            last = _text(au, "LastName")
            fore = _text(au, "ForeName")
            if last:
                authors.append(f"{fore} {last}".strip())

    # Journal
    journal = ""
    if article is not None:
        journal_el = article.find("Journal/Title")
        if journal_el is not None and journal_el.text:
            journal = journal_el.text.strip()
        else:
            journal = _text(article, "Journal/ISOAbbreviation")

    # Date — try multiple locations
    date_str = ""
    date_node = None
    if article is not None:
        date_node = article.find("Journal/JournalIssue/PubDate")
    if date_node is not None:
        y = _text(date_node, "Year")
        m = _text(date_node, "Month")
        d = _text(date_node, "Day")
        medline_date = _text(date_node, "MedlineDate")
        if y:
            parts = [y]
            if m:
                parts.append(m)
            if d:
                parts.append(d)
            date_str = " ".join(parts)
        elif medline_date:
            date_str = medline_date

    # DOI
    doi = ""
    if article is not None:
        for eid in article.findall("ELocationID"):
            if eid.get("EIdType") == "doi" and eid.text:
                doi = eid.text.strip()
                break

    # Abstract
    abstract = ""
    if article is not None:
        abs_node = article.find("Abstract")
        if abs_node is not None:
            parts = []
            for at in abs_node.findall("AbstractText"):
                label = at.get("Label", "")
                text = "".join(at.itertext()).strip()
                if label and text:
                    parts.append(f"{label}: {text}")
                elif text:
                    parts.append(text)
            abstract = " ".join(parts)

    return {
        "pmid": pmid,
        "title": title,
        "authors": authors,
        "journal": journal,
        "date": date_str,
        "doi": doi,
        "abstract": abstract,
        "url": f"https://pubmed.ncbi.nlm.nih.gov/{pmid}/",
    }


# ===================================================================
# Formatting helpers
# ===================================================================

def _author_summary(authors: list[str]) -> str:
    """Return 'First … Last' or full list if ≤ 3."""
    if not authors:
        return "Unknown"
    if len(authors) <= 3:
        return ", ".join(authors)
    return f"{authors[0]} … {authors[-1]}"


def format_paper_terminal(paper: dict) -> str:
    """Pretty terminal output for one paper."""
    lines = [
        f"  📄 {paper['title']}",
        f"     👤 {_author_summary(paper['authors'])}",
        f"     📰 {paper['journal']}  •  {paper['date']}",
        f"     🔗 {paper['url']}",
    ]
    return "\n".join(lines)


def format_paper_markdown(paper: dict) -> str:
    """Telegram-ready markdown for one paper."""
    title = paper["title"].replace("[", "\\[").replace("]", "\\]")
    auth = _author_summary(paper["authors"])
    return (
        f"📄 **{title}**\n"
        f"👤 {auth}\n"
        f"📰 _{paper['journal']}_ • {paper['date']}\n"
        f"🔗 [PubMed {paper['pmid']}]({paper['url']})"
    )


def format_results_json(results: list[dict]) -> str:
    """Full JSON dump of watch results."""
    return json.dumps(results, indent=2, ensure_ascii=False)


# ===================================================================
# Commands
# ===================================================================

def cmd_watch(args):
    """Run all watched queries, report NEW papers only."""
    cfg = load_config()
    seen = load_seen()
    queries = cfg["queries"]
    days = args.days

    if not queries:
        print("⚠️  No queries configured. Add one with: pubmed_watcher.py add \"your query\"")
        return

    all_results = []  # [{query, query_id, new_papers: [...]}, ...]
    total_new = 0

    for i, q in enumerate(queries):
        query_text = q["query"]
        query_id = q["id"]

        if i > 0:
            time.sleep(API_DELAY)

        # Search
        try:
            pmids = search_pmids(query_text, days=days)
        except ConnectionError as exc:
            print(f"❌ Error searching \"{query_text}\": {exc}", file=sys.stderr)
            continue

        # Filter to unseen
        new_pmids = [p for p in pmids if p not in seen]

        if not new_pmids:
            all_results.append({"query": query_text, "query_id": query_id, "new_papers": []})
            continue

        time.sleep(API_DELAY)

        # Fetch details
        try:
            papers = fetch_details(new_pmids)
        except ConnectionError as exc:
            print(f"❌ Error fetching details for \"{query_text}\": {exc}", file=sys.stderr)
            continue

        all_results.append({"query": query_text, "query_id": query_id, "new_papers": papers})
        total_new += len(papers)

        # Mark as seen
        now_iso = datetime.now(timezone.utc).isoformat() + "Z"
        for p in papers:
            seen[p["pmid"]] = now_iso

    # Persist seen
    save_seen(seen)

    # Record history
    history = load_history()
    history_entry = {
        "timestamp": datetime.now(timezone.utc).isoformat() + "Z",
        "days": days,
        "total_new": total_new,
        "queries_run": len(queries),
        "breakdown": [
            {"query": r["query"], "new_count": len(r["new_papers"])}
            for r in all_results
        ],
    }
    history.append(history_entry)
    # Keep last 100 entries
    if len(history) > 100:
        history = history[-100:]
    save_history(history)

    # Output
    if args.json:
        print(format_results_json(all_results))
        return

    if args.markdown:
        _output_markdown(all_results, total_new, days)
        return

    _output_terminal(all_results, total_new, days)


def _output_terminal(results: list[dict], total_new: int, days: int):
    """Pretty-print results to terminal."""
    queries_with_results = [r for r in results if r["new_papers"]]

    if total_new == 0:
        print(f"✅ No new papers found across {len(results)} queries (last {days} days)")
        return

    print(f"\n🔔 {total_new} new paper{'s' if total_new != 1 else ''} found across {len(results)} queries (last {days} days)\n")

    for r in results:
        papers = r["new_papers"]
        if not papers:
            continue
        print(f"━━━ 🔎 \"{r['query']}\" — {len(papers)} new ━━━\n")
        for paper in papers:
            print(format_paper_terminal(paper))
            print()

    print(f"{'━' * 50}")
    print(f"📊 Total: {total_new} new | {len(queries_with_results)}/{len(results)} queries had results")


def _output_markdown(results: list[dict], total_new: int, days: int):
    """Telegram-ready markdown output."""
    if total_new == 0:
        print(f"✅ No new papers found across {len(results)} queries (last {days} days)")
        return

    lines = [
        f"🔔 **{total_new} new paper{'s' if total_new != 1 else ''}** across {len(results)} queries (last {days} days)\n"
    ]

    for r in results:
        papers = r["new_papers"]
        if not papers:
            continue
        lines.append(f"**🔎 {r['query']}** — {len(papers)} new\n")
        for paper in papers:
            lines.append(format_paper_markdown(paper))
            lines.append("")

    print("\n".join(lines))


def cmd_add(args):
    """Add a new watch query."""
    query = args.query.strip()
    if not query:
        print("❌ Query cannot be empty.", file=sys.stderr)
        sys.exit(1)

    cfg = load_config()

    # Check for duplicates
    existing = [q["query"].lower() for q in cfg["queries"]]
    if query.lower() in existing:
        print(f"⚠️  Query already exists: \"{query}\"")
        return

    new_id = cfg["next_id"]
    cfg["queries"].append({
        "id": new_id,
        "query": query,
        "added": datetime.now(timezone.utc).isoformat() + "Z",
    })
    cfg["next_id"] = new_id + 1
    save_config(cfg)
    print(f"✅ Added query #{new_id}: \"{query}\"")


def cmd_list(args):
    """List all configured watch queries."""
    cfg = load_config()
    queries = cfg["queries"]

    if not queries:
        print("📭 No queries configured.")
        return

    print(f"\n📋 Watched Queries ({len(queries)})\n")
    for q in queries:
        added = q.get("added", "?")
        if added and added != "?":
            try:
                dt = datetime.fromisoformat(added.rstrip("Z"))
                added = dt.strftime("%Y-%m-%d")
            except ValueError:
                pass
        print(f"  [{q['id']}] {q['query']}")
        print(f"      Added: {added}")
    print()


def cmd_remove(args):
    """Remove a watch query by ID."""
    cfg = load_config()
    target_id = args.id
    original_len = len(cfg["queries"])
    cfg["queries"] = [q for q in cfg["queries"] if q["id"] != target_id]

    if len(cfg["queries"]) == original_len:
        print(f"❌ No query with ID {target_id} found.", file=sys.stderr)
        sys.exit(1)

    save_config(cfg)
    print(f"🗑️  Removed query #{target_id}")


def cmd_history(args):
    """Show recent alert history."""
    history = load_history()
    limit = args.limit

    if not history:
        print("📭 No alert history yet. Run `watch` to start tracking.")
        return

    recent = history[-limit:]
    print(f"\n📜 Alert History (last {len(recent)} runs)\n")

    for entry in reversed(recent):
        ts = entry.get("timestamp", "?")
        try:
            dt = datetime.fromisoformat(ts.rstrip("Z"))
            ts_display = dt.strftime("%Y-%m-%d %H:%M UTC")
        except ValueError:
            ts_display = ts

        total = entry.get("total_new", 0)
        queries_run = entry.get("queries_run", 0)
        days = entry.get("days", 7)

        icon = "🔔" if total > 0 else "✅"
        print(f"  {icon} {ts_display} — {total} new papers ({queries_run} queries, {days}d window)")

        breakdown = entry.get("breakdown", [])
        for b in breakdown:
            if b["new_count"] > 0:
                print(f"      • {b['query']}: {b['new_count']} new")
    print()


def cmd_reset_seen(args):
    """Reset the seen-papers database (for re-alerting)."""
    if SEEN_PATH.exists():
        backup = SEEN_PATH.with_suffix(".bak.json")
        import shutil
        shutil.copy2(SEEN_PATH, backup)
        print(f"📦 Backed up to {backup.name}")

    save_seen({})
    print("🔄 Seen-papers database reset. Next `watch` will report all recent papers.")


# ===================================================================
# CLI
# ===================================================================

def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        prog="pubmed_watcher",
        description="🔬 PubMed Watcher — Automated alerting for new research papers",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            "  %(prog)s watch                      Run all queries, report new papers\n"
            "  %(prog)s watch --days 14 --markdown  Last 14 days, markdown output\n"
            "  %(prog)s add \"CRISPR aging therapy\"   Add a new query\n"
            "  %(prog)s list                        Show watched queries\n"
            "  %(prog)s remove 3                    Remove query #3\n"
            "  %(prog)s history                     Show recent runs\n"
        ),
    )
    sub = parser.add_subparsers(dest="command", help="Command to run")

    # watch
    p_watch = sub.add_parser("watch", help="Run all queries, report new papers")
    p_watch.add_argument("--days", type=int, default=7, help="Look-back window in days (default: 7)")
    p_watch.add_argument("--json", action="store_true", help="Machine-readable JSON output")
    p_watch.add_argument("--markdown", action="store_true", help="Telegram-ready markdown output")
    p_watch.set_defaults(func=cmd_watch)

    # add
    p_add = sub.add_parser("add", help="Add a new watch query")
    p_add.add_argument("query", type=str, help="PubMed search query to watch")
    p_add.set_defaults(func=cmd_add)

    # list
    p_list = sub.add_parser("list", help="Show all watched queries")
    p_list.set_defaults(func=cmd_list)

    # remove
    p_remove = sub.add_parser("remove", help="Remove a watch query by ID")
    p_remove.add_argument("id", type=int, help="Query ID to remove")
    p_remove.set_defaults(func=cmd_remove)

    # history
    p_history = sub.add_parser("history", help="Show recent alert history")
    p_history.add_argument("--limit", type=int, default=20, help="Number of entries to show (default: 20)")
    p_history.set_defaults(func=cmd_history)

    # reset-seen
    p_reset = sub.add_parser("reset-seen", help="Clear seen-papers DB (re-alert everything)")
    p_reset.set_defaults(func=cmd_reset_seen)

    return parser


def main():
    parser = build_parser()
    args = parser.parse_args()

    if not args.command:
        parser.print_help()
        sys.exit(0)

    args.func(args)


if __name__ == "__main__":
    main()