app.py

"""Google Scholar citation tracker — FastAPI + Bootstrap SPA on Hugging Face (Docker).

Backend serves a static Bootstrap page (index.html) plus JSON:
  GET /                     -> the single-page app
  GET /api/data             -> everything the page needs (author, totals, history, meta)
  GET /badge/total          -> shields.io endpoint JSON (total citations, with icon)
  GET /badge/paper/{id}     -> shields.io endpoint JSON for one paper (stable id)
  GET /api/summary          -> compact state (kept for keep-alive pings)

A daily job records one snapshot per day into a HF Dataset repo. On restart the same
day it skips the SerpApi call (quota guard).

Env: SCHOLAR_URL (req), SERPAPI_KEY (req on HF), HF_TOKEN (req), HF_DATASET_REPO (opt),
     DAILY_HOUR (opt, UTC), SPACE_HOST / SPACE_ID (auto on HF).
"""

import os
import logging
import threading
from datetime import date
from urllib.parse import quote

import pandas as pd
from fastapi import FastAPI
from fastapi.responses import JSONResponse, FileResponse

import scholar
from storage import Storage, TOTAL_ID

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("app")

HERE = os.path.dirname(os.path.abspath(__file__))

# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #
SCHOLAR_URL = os.environ.get("SCHOLAR_URL", "").strip()
HF_TOKEN = os.environ.get("HF_TOKEN")
DAILY_HOUR = int(os.environ.get("DAILY_HOUR", "3"))
SPACE_HOST = os.environ.get("SPACE_HOST", "").strip()
# Collapse the dataset repo's git history to a single commit after each daily
# write, so accumulated commit blobs never bloat the repo. Data lives in the CSV
# files, not git history, so this is lossless. Set DATASET_SQUASH=0 to disable.
DATASET_SQUASH = os.environ.get("DATASET_SQUASH", "1").strip() not in ("0", "false", "False", "")


def _derived_dataset_repo() -> str | None:
    space_id = os.environ.get("SPACE_ID")
    if not space_id or "/" not in space_id:
        return None
    owner, name = space_id.split("/", 1)
    return f"{owner}/{name}-data"


def _writable_namespaces() -> set[str]:
    """HF namespaces our token can push to (the user itself + their orgs)."""
    if not HF_TOKEN:
        return set()
    try:
        from huggingface_hub import HfApi

        me = HfApi(token=HF_TOKEN).whoami()
        ns = {me.get("name")}
        for org in me.get("orgs", []) or []:
            ns.add(org.get("name"))
        return {n for n in ns if n}
    except Exception:  # noqa: BLE001
        return set()


def _resolve_dataset_repo() -> str | None:
    """Pick the dataset repo, self-healing an inherited / non-writable value.

    When a Space is duplicated, HF copies the ORIGINAL's Variables (values included)
    into the fork — so the fork starts with the original owner's HF_DATASET_REPO. The
    forker can't write there, so we detect it (owner not in our writable namespaces)
    and fall back to our OWN auto-derived {owner}/{space}-data, then re-pin the
    corrected value. Net effect: forkers do nothing; renames within your own account
    stay linked to the same dataset.
    """
    env = (os.environ.get("HF_DATASET_REPO") or "").strip() or None
    derived = _derived_dataset_repo()
    writable = _writable_namespaces()

    # Respect an explicit value only if it lives in a namespace we can push to.
    # (If we couldn't determine namespaces, don't second-guess the user.)
    if env and (not writable or env.split("/", 1)[0] in writable):
        chosen = env
    else:
        chosen = derived or env

    # Persist the resolved value as a Space Variable (best-effort), unless unchanged.
    space_id = os.environ.get("SPACE_ID")
    if chosen and chosen != env and space_id and HF_TOKEN:
        try:
            from huggingface_hub import HfApi

            HfApi(token=HF_TOKEN).add_space_variable(
                repo_id=space_id, key="HF_DATASET_REPO", value=chosen
            )
            log.info("Resolved HF_DATASET_REPO=%s (was %r).", chosen, env)
        except Exception:  # noqa: BLE001
            log.exception("Could not pin HF_DATASET_REPO (non-fatal).")
    return chosen


HF_DATASET_REPO = _resolve_dataset_repo()
storage = Storage(HF_DATASET_REPO, HF_TOKEN)

STATE = {
    "name": "", "affiliation": "", "total": 0, "updated": None,
    "papers": [], "papers_by_id": {}, "error": None,
}
_STATE_LOCK = threading.Lock()


# --------------------------------------------------------------------------- #
# Core: fetch + record
# --------------------------------------------------------------------------- #
def _refresh_state_from_df(df: pd.DataFrame) -> None:
    if df.empty:
        return
    latest_day = df["date"].max()
    latest = df[df["date"] == latest_day]

    papers, by_id, total = [], {}, 0
    name = aff = None
    for _, row in latest.iterrows():
        if row["paper_id"] == TOTAL_ID:
            total = int(row["citations"])
            # Author name/affiliation are persisted in the TOTAL row (title/year).
            if row["title"] and row["title"] != "All Papers":
                name = row["title"]
            if row["year"]:
                aff = row["year"]
            continue
        papers.append({"paper_id": row["paper_id"], "title": row["title"],
                       "year": row["year"], "venue": row.get("venue", ""),
                       "citations": int(row["citations"])})
        by_id[row["paper_id"]] = int(row["citations"])
    papers.sort(key=lambda p: p["citations"], reverse=True)
    with _STATE_LOCK:
        STATE["papers"] = papers
        STATE["papers_by_id"] = by_id
        STATE["total"] = total
        STATE["updated"] = latest_day
        if name is not None:
            STATE["name"] = name
        if aff is not None:
            STATE["affiliation"] = aff


def load_state():
    try:
        _refresh_state_from_df(storage.load())
    except Exception:  # noqa: BLE001
        log.exception("load_state failed")


def record_today(force: bool = False) -> str:
    if not SCHOLAR_URL:
        return "❌ SCHOLAR_URL is not set."

    today = str(date.today())
    df = storage.load()

    if not force and not df.empty and (df["date"] == today).any():
        _refresh_state_from_df(df)
        log.info("Today (%s) already recorded — skipping fetch (SerpApi quota saved).", today)
        return f"✅ Today ({today}) already recorded."

    try:
        data = scholar.fetch_author(SCHOLAR_URL)
    except Exception as e:  # noqa: BLE001
        log.exception("Fetch failed")
        with _STATE_LOCK:
            STATE["error"] = str(e)
        _refresh_state_from_df(df)
        return f"⚠️ Fetch failed ({scholar.active_backend()}): {e}."

    storage.write_snapshot(
        today, data["scholar_id"], data["name"], data["affiliation"],
        data["total_citations"], data["publications"],
    )
    if DATASET_SQUASH:
        storage.squash_history()

    df = storage.load()
    with _STATE_LOCK:
        STATE["name"] = data["name"]
        STATE["affiliation"] = data["affiliation"]
        STATE["error"] = None
    _refresh_state_from_df(df)
    return f"✅ Recorded {len(data['publications'])} papers. Total: {data['total_citations']}."


# --------------------------------------------------------------------------- #
# Helpers
# --------------------------------------------------------------------------- #
def _host():
    return SPACE_HOST or "YOUR-SPACE-HOST.hf.space"


def _shield_url(endpoint_path: str) -> str:
    inner = f"https://{_host()}{endpoint_path}"
    return f"https://img.shields.io/endpoint?url={quote(inner, safe='')}"


def _dataset_size_str() -> str:
    if not storage.remote:
        return "local CSV"
    try:
        info = storage.api.repo_info(HF_DATASET_REPO, repo_type="dataset", files_metadata=True)
        size = float(sum((s.size or 0) for s in info.siblings))
        for unit in ["B", "KB", "MB", "GB"]:
            if size < 1024:
                return f"{size:.1f} {unit}"
            size /= 1024
        return f"{size:.1f} TB"
    except Exception as e:  # noqa: BLE001
        return f"unknown"


def build_data() -> dict:
    df = storage.load()
    _refresh_state_from_df(df)
    with _STATE_LOCK:
        name, aff = STATE["name"], STATE["affiliation"]
        total, updated, err = STATE["total"], STATE["updated"], STATE["error"]
        papers = list(STATE["papers"])

    dates = sorted(df["date"].dropna().astype(str).unique().tolist())

    def series(mask):
        s = df[mask].set_index("date")["citations"]
        return [int(s[d]) if d in s.index else None for d in dates]

    history = {
        "dates": dates,
        "total": series(df["paper_id"] == TOTAL_ID),
        "by_paper": {p["paper_id"]: series(df["paper_id"] == p["paper_id"]) for p in papers},
    }
    today = str(date.today())
    return {
        "author": {"name": name, "affiliation": aff, "scholar_url": SCHOLAR_URL},
        "total": total,
        "updated": updated,
        "fetched_today": updated == today,
        "today": today,
        "daily_hour": DAILY_HOUR,
        "space_host": _host(),
        "dataset_repo": HF_DATASET_REPO or "",
        "dataset_size": _dataset_size_str(),
        "backend": scholar.active_backend(),
        "error": err,
        "papers": papers,
        "history": history,
        "badge_base": f"https://{_host()}",
    }


# --------------------------------------------------------------------------- #
# FastAPI
# --------------------------------------------------------------------------- #
app = FastAPI(title="Paper Citation Tracker")


def _badge(message, fmt="full", color="blue"):
    # fmt controls the three display styles the UI offers:
    #   "num"     -> number only (no icon, no label)
    #   "iconnum" -> Scholar icon + number (no label)
    #   "full"    -> Scholar icon + "Citations" label + number (default)
    payload = {
        "schemaVersion": 1,
        "label": "Citations" if fmt == "full" else "",
        "message": str(message),
        "color": color,
        "cacheSeconds": 300,
    }
    if fmt in ("iconnum", "full"):
        payload["namedLogo"] = "googlescholar"
        payload["logoColor"] = "white"
    return JSONResponse(
        payload,
        headers={"Cache-Control": "no-cache, max-age=0"},
    )


_NOCACHE = {"Cache-Control": "no-cache, max-age=0"}


@app.get("/")
def index():
    return FileResponse(os.path.join(HERE, "index.html"), headers=_NOCACHE)


@app.get("/embed")
def embed():
    return FileResponse(os.path.join(HERE, "embed.html"), headers=_NOCACHE)


@app.get("/api/data")
def api_data():
    return JSONResponse(build_data())


@app.get("/api/summary")
def api_summary():
    with _STATE_LOCK:
        return {k: v for k, v in STATE.items() if k != "papers_by_id"}


@app.get("/badge/total")
def badge_total(fmt: str = "full"):
    with _STATE_LOCK:
        return _badge(STATE["total"], fmt=fmt)


@app.get("/badge/paper/{paper_id}")
def badge_paper(paper_id: str, fmt: str = "full"):
    with _STATE_LOCK:
        c = STATE["papers_by_id"].get(paper_id)
    return _badge("n/a", fmt=fmt, color="lightgrey") if c is None else _badge(c, fmt=fmt)


# --------------------------------------------------------------------------- #
# Startup + daily schedule
# --------------------------------------------------------------------------- #
def _startup():
    load_state()
    try:
        log.info(record_today(force=False))
    except Exception:  # noqa: BLE001
        log.exception("Startup record failed")


def _schedule():
    try:
        from apscheduler.schedulers.background import BackgroundScheduler

        sched = BackgroundScheduler(timezone="UTC")
        sched.add_job(lambda: log.info(record_today(force=False)), "cron", hour=DAILY_HOUR, minute=0)
        sched.start()
        log.info("Daily job scheduled at %02d:00 UTC.", DAILY_HOUR)
    except Exception:  # noqa: BLE001
        log.exception("Scheduler setup failed")


threading.Thread(target=_startup, daemon=True).start()
_schedule()


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=int(os.environ.get("PORT", "7860")))