services/storage.py

"""
Everything related to talking to the Hugging Face dataset that acts as
our "database": checking whether a plot was already generated, queueing a
new Plotly figure for the next batched upload, and keeping a small
catalog/index file per PDF set so lookups don't require listing the whole
repo.

Plots are NOT uploaded one at a time. Hugging Face caps commits at
128/hour per repo, and the old immediate-upload path spent two commits
per plot (the plot file + the index.json update) — a handful of dashboard
visitors could exhaust that budget in minutes. Instead, `queue_plot`
writes the figure to local scratch space and registers it in memory;
a background job (started in app.py) periodically calls `flush_pending`,
which bundles every queued plot across every grid into a SINGLE
`create_commit` call. Until a flush happens, generated plots are served
directly off local disk by app.py's `/plots/{folder}/{filename}` route,
so callers see the plot immediately regardless of commit timing.

Dataset layout (inside `dataset_repo`):

    <grid_name>/                          <- LHAPDF grid files (already existing)
    <grid_name>plots/
        index.json                         <- catalog of every plot generated for this grid
        <grid_name>_plt_xfx_PID<id>_Q<q>.json               <- standard plot
        <grid_name>_plt_xfx_ratio_PID<id>_Q<q>.json          <- ratio plot
        <grid_name>_plt_xfx_corr_PID<id1>_PID<id2>_Q<q>.json <- correlation plot
"""

import hashlib
import json
import logging
import os
import threading
import time
from datetime import datetime, timezone
from typing import Optional

from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download
from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

from services.config import SPACE_BASE_URL, TMP_DIR, index_path_in_repo, plots_folder

logger = logging.getLogger(__name__)

_lock = threading.Lock()
# plot_id -> index entry, for plots generated but not yet committed to HF.
_pending_entries: dict = {}
# "dataset_repo|grid_name" -> list of plot_ids queued for that grid, so a
# flush can group per-grid index.json updates without scanning everything.
_pending_by_grid: dict = {}


def _grid_key(dataset_repo: str, grid_name: str) -> str:
    return f"{dataset_repo}|{grid_name}"


def make_plot_id(plot_type: str, grid_name: str, parameters: dict) -> str:
    """
    Deterministic id for a plot request: same plot_type + grid_name +
    parameters always hashes to the same id, which is how we detect an
    already-generated plot regardless of key ordering. Used only for
    cache lookups in index.json — the actual filename stored in the
    dataset follows the human-readable convention in services.config.
    """
    payload = json.dumps(
        {"plot_type": plot_type, "grid_name": grid_name, "parameters": parameters},
        sort_keys=True,
    )
    digest = hashlib.sha256(payload.encode("utf-8")).hexdigest()[:16]
    return f"{plot_type}_{digest}"


def _local_plot_path(path_in_repo: str) -> str:
    return os.path.join(TMP_DIR, path_in_repo)


def _load_index(dataset_repo: str, grid_name: str, hf_token: str) -> list:
    """Download and parse index.json for a grid. Returns [] if it doesn't exist yet."""
    try:
        local_path = hf_hub_download(
            repo_id=dataset_repo,
            repo_type="dataset",
            filename=index_path_in_repo(grid_name),
            token=hf_token,
        )
    except (EntryNotFoundError, RepositoryNotFoundError, Exception):
        # Covers "file not found" as well as "folder doesn't exist yet".
        return []

    with open(local_path, "r") as f:
        return json.load(f)


def find_cached_plot(dataset_repo: str, grid_name: str, plot_id: str, hf_token: str) -> Optional[dict]:
    """Look up a plot that's already been committed to the HF dataset's index.json."""
    index = _load_index(dataset_repo, grid_name, hf_token)
    for entry in index:
        if entry.get("plot_id") == plot_id:
            logger.info("plot_id=%s found in committed index.json for %s/%s", plot_id, dataset_repo, grid_name)
            return entry
    logger.info("plot_id=%s not in committed index.json for %s/%s (%d entries checked)", plot_id, dataset_repo, grid_name, len(index))
    return None


def find_pending_plot(plot_id: str) -> Optional[dict]:
    """
    Look up a plot generated earlier in the current batching window but
    not yet flushed to HF. Checked before recomputing so repeat requests
    for the same plot (e.g. a page reload) don't waste compute or queue
    duplicate uploads while waiting for the next flush.
    """
    with _lock:
        entry = _pending_entries.get(plot_id)
    logger.info("plot_id=%s %s in pending (not-yet-flushed) entries", plot_id, "found" if entry else "not found")
    return entry


def queue_plot(
    dataset_repo: str,
    grid_name: str,
    plot_id: str,
    plot_type: str,
    parameters: dict,
    fig,
    json_filename: str,
) -> dict:
    """
    Write a Plotly figure to local scratch space and register it for the
    next batched commit (see flush_pending). Returns the same shape as
    the old immediate-upload path so callers don't need to change, except
    plotly_json_url points at this Space's own /plots route (served
    straight off local disk) rather than the HF resolve URL, since the
    file may not be committed yet.
    """
    path_in_repo = f"{plots_folder(grid_name)}/{json_filename}"
    local_path = _local_plot_path(path_in_repo)
    os.makedirs(os.path.dirname(local_path), exist_ok=True)
    with open(local_path, "w") as f:
        f.write(fig.to_json())

    entry = {
        "plot_id": plot_id,
        "plot_type": plot_type,
        "grid_name": grid_name,
        "parameters": parameters,
        "json_filename": json_filename,
        "json_path": path_in_repo,
        "dataset_repo": dataset_repo,
        "plot_url": f"https://huggingface.co/datasets/{dataset_repo}/blob/main/{path_in_repo}",
        "plotly_json_url": f"{SPACE_BASE_URL}/plots/{path_in_repo}",
        "created_at": datetime.now(timezone.utc).isoformat(),
    }

    key = _grid_key(dataset_repo, grid_name)
    with _lock:
        _pending_entries[plot_id] = entry
        _pending_by_grid.setdefault(key, [])
        if plot_id not in _pending_by_grid[key]:
            _pending_by_grid[key].append(plot_id)
        pending_count = len(_pending_by_grid[key])

    logger.info("queued plot_id=%s (%s) for %s -- %d plot(s) pending flush for this grid", plot_id, json_filename, key, pending_count)
    return entry


def pending_status() -> dict:
    """Number of not-yet-flushed plots queued per grid, for the /pending-status endpoint."""
    with _lock:
        return {key: len(ids) for key, ids in _pending_by_grid.items()}


def flush_pending(hf_token: str) -> dict:
    """
    Bundle every queued plot (across every grid/dataset_repo) into ONE
    `create_commit` call per dataset_repo: all plot JSON files plus each
    affected grid's merged index.json. This is what keeps the Space
    under HF's per-repo commit rate limit regardless of how many plots
    were generated during the window.
    """
    with _lock:
        if not _pending_entries:
            logger.info("flush_pending: nothing queued, skipping")
            return {"flushed": 0, "grids": []}
        entries_by_plot_id = dict(_pending_entries)
        by_grid = {key: list(ids) for key, ids in _pending_by_grid.items()}

    logger.info("flush_pending: starting flush of %d plot(s) across %d grid(s)", len(entries_by_plot_id), len(by_grid))

    # Group by dataset_repo since create_commit operates on a single repo.
    operations_by_repo: dict = {}
    flushed_plot_ids = []
    flushed_grid_keys = []

    for key, plot_ids in by_grid.items():
        dataset_repo, grid_name = key.split("|", 1)
        pending_for_grid = [entries_by_plot_id[pid] for pid in plot_ids if pid in entries_by_plot_id]
        if not pending_for_grid:
            continue

        operations = operations_by_repo.setdefault(dataset_repo, [])

        for entry in pending_for_grid:
            local_path = _local_plot_path(entry["json_path"])
            if os.path.isfile(local_path):
                operations.append(
                    CommitOperationAdd(path_in_repo=entry["json_path"], path_or_fileobj=local_path)
                )

        existing_index = _load_index(dataset_repo, grid_name, hf_token)
        merged = [e for e in existing_index if e.get("plot_id") not in plot_ids]
        merged.extend(
            {k: v for k, v in e.items() if k != "dataset_repo"} for e in pending_for_grid
        )
        local_index_path = os.path.join(TMP_DIR, f"_flush_index_{dataset_repo.replace('/', '_')}_{grid_name}.json")
        with open(local_index_path, "w") as f:
            json.dump(merged, f, indent=2)
        operations.append(
            CommitOperationAdd(path_in_repo=index_path_in_repo(grid_name), path_or_fileobj=local_index_path)
        )

        flushed_plot_ids.extend(plot_ids)
        flushed_grid_keys.append(key)

    if not any(operations_by_repo.values()):
        logger.info("flush_pending: no operations to commit after grouping, skipping")
        return {"flushed": 0, "grids": []}

    api = HfApi(token=hf_token)
    for dataset_repo, operations in operations_by_repo.items():
        if not operations:
            continue
        logger.info("flush_pending: committing %d operation(s) to %s ...", len(operations), dataset_repo)
        t0 = time.monotonic()
        api.create_commit(
            repo_id=dataset_repo,
            repo_type="dataset",
            operations=operations,
            commit_message=f"Batch upload of {len(operations)} file(s)",
        )
        logger.info("flush_pending: committed to %s in %.2fs", dataset_repo, time.monotonic() - t0)

    # Only after every commit above has actually succeeded: delete the local
    # copies of everything just uploaded. Local disk on the Space container
    # is scratch space, not the source of truth (the HF dataset is), so
    # there's no reason to keep these around once they're safely committed —
    # they'd otherwise accumulate forever. Plots not yet flushed are left
    # alone (still needed for local-serving and repeat-request dedup); the
    # per-grid temp index snapshot files are always safe to remove, flushed
    # or not, since they're rebuilt fresh on every flush.
    for key in flushed_grid_keys:
        for pid in by_grid.get(key, []):
            entry = entries_by_plot_id.get(pid)
            if not entry:
                continue
            local_path = _local_plot_path(entry["json_path"])
            try:
                os.remove(local_path)
            except OSError:
                pass  # already gone, or never existed — nothing to clean up

    for dataset_repo in operations_by_repo:
        for key in flushed_grid_keys:
            repo_of_key, grid_name = key.split("|", 1)
            if repo_of_key != dataset_repo:
                continue
            local_index_path = os.path.join(TMP_DIR, f"_flush_index_{dataset_repo.replace('/', '_')}_{grid_name}.json")
            try:
                os.remove(local_index_path)
            except OSError:
                pass

    with _lock:
        for pid in flushed_plot_ids:
            _pending_entries.pop(pid, None)
        for key in flushed_grid_keys:
            remaining = [pid for pid in _pending_by_grid.get(key, []) if pid not in flushed_plot_ids]
            if remaining:
                _pending_by_grid[key] = remaining
            else:
                _pending_by_grid.pop(key, None)

    logger.info("flush_pending: done -- flushed %d plot(s) across %d grid(s)", len(flushed_plot_ids), len(flushed_grid_keys))

    return {"flushed": len(flushed_plot_ids), "grids": flushed_grid_keys}