server/persistence.py

"""Menagerie persistence + HF dataset sync + the pending-record vault.

Durable truth is ``menagerie/``: one ``records.jsonl`` line per published
awakening plus its media siblings ``{id}.jpg`` / ``{id}.wav``. Record ids are
server-generated tokens; the media route serves ONLY through the id->path map
built here, so no client-supplied string ever becomes a filesystem path.
Records carry NO user identity — privacy is part of the contract.

Dataset sync (godseed pattern, batched): when ``PAREIDOLIA_DATASET`` is set
(token from ``HF_TOKEN`` or the logged-in CLI), dirty files are pushed as ONE
``create_commit`` of ``CommitOperationAdd`` ops, debounced to at most one
commit per 60s. Every network error is swallowed and logged — the wall must
never die of a flaky connection; failed batches retry on the next publish.

Boot restore order (ARCHITECTURE.md §5): HF dataset -> seeds/records/ -> empty.

``PendingStore`` is the ephemeral half: awaken parks the full record (+ media
bytes) here under a one-time ``record_token`` so the publish step cannot be
forged or tampered client-side. 15-minute expiry, single use, bounded size.
"""

from __future__ import annotations

import asyncio
import json
import logging
import os
import re
import secrets
import shutil
import threading
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Optional

log = logging.getLogger("pareidolia.persistence")

RECORDS_FILENAME = "records.jsonl"
MEDIA_KINDS = ("jpg", "wav")
# Optional media siblings beyond {id}.jpg/{id}.wav: the audible idle mutter
# and the share-page OG composite. Registered under their full stem (e.g.
# ("m_abc_mutter", "wav")) so the /media route's stem.ext parse finds them
# through the same id->path map — no second lookup scheme.
MEDIA_SIBLINGS = (("_mutter", "wav"), ("_og", "jpg"))
DEBOUNCE_SECONDS = 60.0
PENDING_TTL_SECONDS = 15 * 60.0
MAX_PENDING = 256  # each pending holds image+wav bytes; bound the unpublished pile

# Record/media ids must be exactly what we mint (token_urlsafe alphabet). The
# /media route's regex mirrors this; anything else never reaches the map.
RECORD_ID_RE = re.compile(r"^[A-Za-z0-9_-]{4,64}$")


def new_record_id() -> str:
    """Server-generated public id (doubles as the media filename stem)."""
    return f"m_{secrets.token_urlsafe(9)}"


def _safe_exc(exc: BaseException, secret: Optional[str] = None) -> str:
    """One-line, token-safe rendering of an exception (security review #4).

    The dataset-sync call sites hold the HF token in scope; logging raw
    exception objects there risks echoing it (hub errors can embed request
    details). Log the exception TYPE plus a truncated single-line message with
    the secret scrubbed if it appears — never the raw object."""
    text = str(exc).replace("\n", " ").replace("\r", " ")
    if secret:
        text = text.replace(secret, "[redacted]")
    return f"{type(exc).__name__}: {text[:200]}"


# --------------------------------------------------------------------------- store
class MenagerieStore:
    """Append-only JSONL + media files under ``menagerie/``. Thread-safe appends
    (uploads run in worker threads; the boot warm-up may too)."""

    def __init__(self, menagerie_dir: Path | str) -> None:
        self.dir = Path(menagerie_dir)
        self.dir.mkdir(parents=True, exist_ok=True)
        self.path = self.dir / RECORDS_FILENAME
        self._lock = threading.Lock()

    def append(self, record: dict[str, Any]) -> dict[str, Any]:
        encoded = json.dumps(record, ensure_ascii=False)
        with self._lock:
            with self.path.open("a", encoding="utf-8") as fh:
                fh.write(encoded + "\n")
        return record

    def load(self) -> list[dict[str, Any]]:
        """All record lines, in publish order. Corrupt lines are skipped with a
        warning (a half-written final line after a crash must not brick boot)."""
        if not self.path.exists():
            return []
        records: list[dict[str, Any]] = []
        with self.path.open("r", encoding="utf-8") as fh:
            for lineno, raw in enumerate(fh, start=1):
                raw = raw.strip()
                if not raw:
                    continue
                try:
                    obj = json.loads(raw)
                except json.JSONDecodeError:
                    log.warning("skipping corrupt record line %d in %s", lineno, self.path)
                    continue
                if isinstance(obj, dict):
                    records.append(obj)
        return records

    def write_media(
        self, record_id: str, image_jpeg: bytes, grudge_wav: bytes
    ) -> tuple[Path, Path]:
        """Write the two media siblings for ``record_id``; returns their paths."""
        jpg = self.dir / f"{record_id}.jpg"
        wav = self.dir / f"{record_id}.wav"
        jpg.write_bytes(image_jpeg)
        wav.write_bytes(grudge_wav)
        return jpg, wav

    def write_blob(self, filename: str, data: bytes) -> Path:
        """Write one extra media sibling (mutter wav / OG composite)."""
        path = self.dir / filename
        path.write_bytes(data)
        return path


# --------------------------------------------------------------------- dataset sync
class DatasetSync:
    """Mirror menagerie/ to/from an HF dataset repo (Sharing is Caring evidence
    + survival across Space restarts). Inert unless ``PAREIDOLIA_DATASET`` is
    set; the token falls back from ``HF_TOKEN`` to the logged-in CLI (None lets
    huggingface_hub resolve it). Never raises for network problems."""

    def __init__(
        self,
        store: MenagerieStore,
        token: Optional[str] = None,
        dataset: Optional[str] = None,
    ) -> None:
        self.store = store
        self.token = token if token is not None else os.environ.get("HF_TOKEN")
        self.dataset = (
            dataset if dataset is not None else os.environ.get("PAREIDOLIA_DATASET")
        )
        self._repo_ensured = False

    @property
    def enabled(self) -> bool:
        return bool(self.dataset)

    def restore(self) -> bool:
        """Pull records.jsonl + media from the dataset into the local menagerie
        (fresh Space boot). Returns True only if records.jsonl now exists."""
        if not self.enabled or self.store.path.exists():
            return self.store.path.exists()
        try:
            from huggingface_hub import snapshot_download

            snapshot_download(
                repo_id=self.dataset,
                repo_type="dataset",
                token=self.token,
                local_dir=str(self.store.dir),
                allow_patterns=[RECORDS_FILENAME, "*.jpg", "*.wav"],
            )
            restored = self.store.path.exists()
            if restored:
                log.info("menagerie restored from dataset %s", self.dataset)
            return restored
        except Exception as exc:  # noqa: BLE001 — sync must never crash the app
            log.warning(
                "dataset restore skipped (%s): %s",
                self.dataset,
                _safe_exc(exc, self.token),
            )
            return False

    def _ensure_repo(self) -> None:
        if self._repo_ensured:
            return
        try:
            from huggingface_hub import create_repo

            create_repo(
                self.dataset, repo_type="dataset", token=self.token, exist_ok=True
            )
        except Exception as exc:  # noqa: BLE001 — create_commit will tell us anyway
            log.debug("create_repo skipped: %s", _safe_exc(exc, self.token))
        self._repo_ensured = True

    def upload_batch(self, paths: list[Path]) -> bool:
        """Push ``paths`` as ONE commit (CommitOperationAdd per file). Returns
        success; the caller re-queues the batch on failure."""
        if not self.enabled:
            return False
        existing = [p for p in paths if p.is_file()]
        if not existing:
            return True
        try:
            from huggingface_hub import CommitOperationAdd, HfApi

            self._ensure_repo()
            operations = [
                CommitOperationAdd(path_in_repo=p.name, path_or_fileobj=str(p))
                for p in existing
            ]
            HfApi(token=self.token).create_commit(
                repo_id=self.dataset,
                repo_type="dataset",
                operations=operations,
                commit_message=f"pareidolia: the menagerie grows (+{len(existing)} files)",
            )
            return True
        except Exception as exc:  # noqa: BLE001 — sync must never crash the app
            log.warning(
                "dataset upload failed (next publish retries): %s",
                _safe_exc(exc, self.token),
            )
            return False


# ------------------------------------------------------------------------- service
class PersistenceService:
    """What the rest of the app talks to.

    Boot:    ``service.boot()`` — restore (dataset -> seeds -> empty), load
             records, build the id->path media map. Idempotent + thread-safe;
             afterwards reads come from memory, no file IO per request.
    Publish: ``await service.publish(record, jpeg, wav)`` — mint id, write
             media, append the record line, mark files dirty for the debounced
             dataset batch. Returns the public record (with media URLs).
    """

    def __init__(
        self,
        menagerie_dir: Path | str,
        sync: Optional[DatasetSync] = None,
        seeds_dir: Optional[Path | str] = None,
        debounce: float = DEBOUNCE_SECONDS,
        clock: Callable[[], float] = time.monotonic,
    ) -> None:
        self.store = MenagerieStore(menagerie_dir)
        self.sync = sync if sync is not None else DatasetSync(self.store)
        self.seeds_dir = Path(seeds_dir) if seeds_dir is not None else None
        self.debounce = debounce
        self.clock = clock
        self._records: list[dict[str, Any]] = []
        self._media: dict[tuple[str, str], Path] = {}
        self._booted = False
        self._boot_lock = threading.Lock()
        self._dirty: set[Path] = set()
        self._last_upload: Optional[float] = None
        self._sync_task: Optional[asyncio.Task[Any]] = None

    # ------------------------------------------------------------------------- boot
    @property
    def booted(self) -> bool:
        return self._booted

    def boot(self) -> None:
        """Restore order: HF dataset -> seeds/records/ fallback -> empty."""
        with self._boot_lock:
            if self._booted:
                return
            if not self.store.path.exists():
                if not self.sync.restore():
                    self._restore_from_seeds()
            self._records = self.store.load()
            for record in self._records:
                self._register_media(record)
            self._booted = True
            log.info(
                "menagerie booted: %d records, dataset=%s",
                len(self._records),
                self.sync.dataset if self.sync.enabled else "off",
            )

    def _restore_from_seeds(self) -> None:
        """Copy the committed seed menagerie in (license-safe photos run through
        the real pipeline) so the wall is alive from visitor #1."""
        seeds = self.seeds_dir
        if seeds is None or not (seeds / RECORDS_FILENAME).is_file():
            return
        try:
            for item in sorted(seeds.iterdir()):
                if item.is_file() and (
                    item.name == RECORDS_FILENAME or item.suffix in (".jpg", ".wav")
                ):
                    shutil.copy2(item, self.store.dir / item.name)
            log.info("menagerie seeded from %s", seeds)
        except OSError as exc:
            log.warning("seed restore failed (starting empty): %s", exc)

    def _register_media(self, record: dict[str, Any]) -> None:
        """Admit a record's media into the serve map — ONLY server/seed-written
        ids that match the mint alphabet ever become servable paths.

        Self-healing: a file missing from menagerie/ is recovered from the
        committed seeds/ dir when present. This closes a live June-12 gap —
        the dataset restore can win the boot race holding records.jsonl but
        not every record's media (seed media predated the dataset sync), and
        restore() short-circuits whenever records.jsonl already exists on a
        reused container, so without healing those records 404 forever."""
        record_id = str(record.get("id", ""))
        if not RECORD_ID_RE.fullmatch(record_id):
            return
        for kind in MEDIA_KINDS:
            path = self._present_or_healed(f"{record_id}.{kind}")
            if path is not None:
                self._media[(record_id, kind)] = path
        for suffix, kind in MEDIA_SIBLINGS:
            path = self._present_or_healed(f"{record_id}{suffix}.{kind}")
            if path is not None:
                self._media[(f"{record_id}{suffix}", kind)] = path

    def _present_or_healed(self, filename: str) -> Optional[Path]:
        """The serve path for ``filename`` if it exists locally — else copy it
        in from seeds/ when available there (best-effort), else None."""
        path = self.store.dir / filename
        if path.is_file():
            return path
        seeds = self.seeds_dir
        if seeds is not None and (seeds / filename).is_file():
            try:
                shutil.copy2(seeds / filename, path)
                log.info("media healed from seeds: %s", filename)
                return path
            except OSError as exc:
                log.warning("media heal failed for %s: %s", filename, exc)
        return None

    # ------------------------------------------------------------------------ reads
    def records(self) -> list[dict[str, Any]]:
        if not self._booted:
            self.boot()
        return self._records

    def count(self) -> int:
        return len(self.records())

    def page(self, offset: int, limit: int) -> tuple[list[dict[str, Any]], int]:
        """Newest-first page of public records."""
        newest_first = list(reversed(self.records()))
        return newest_first[offset : offset + limit], len(newest_first)

    def media_path(self, record_id: str, kind: str) -> Optional[Path]:
        """Path-traversal-safe media lookup: the map is the ONLY route from a
        client string to a file; unknown ids and kinds simply don't exist."""
        if kind not in MEDIA_KINDS:
            return None
        return self._media.get((record_id, kind))

    def find(self, record_id: str) -> Optional[dict[str, Any]]:
        """The published record for ``record_id`` (the share page's read), or
        None. Linear scan on purpose: the wall is hundreds of records, the
        list is already in memory, and publish order stays the single truth."""
        for record in self.records():
            if record.get("id") == record_id:
                return record
        return None

    # ---------------------------------------------------------------------- publish
    async def publish(
        self,
        record: dict[str, Any],
        image_jpeg: bytes,
        grudge_wav: bytes,
        mutter_wav: Optional[bytes] = None,
        og_jpeg: Optional[bytes] = None,
    ) -> dict[str, Any]:
        """Persist one awakening to the public wall. The record dict must carry
        no user identity (the contract); this method adds id/ts/media URLs.

        ``mutter_wav`` and ``og_jpeg`` are best-effort siblings (contract):
        the audible idle mutter ({id}_mutter.wav -> ``mutter_audio`` URL on
        the record) and the share-page OG composite ({id}_og.jpg, no record
        key — the /o/ page derives its URL from the id). Either may be None
        (TTS window margin ran out, composite failed) and the publish is
        unchanged — old records and new minimal records stay one shape."""
        if not self._booted:
            await asyncio.to_thread(self.boot)
        record_id = new_record_id()
        public = {
            "id": record_id,
            "ts": time.time(),
            **record,
            "image_url": f"/media/{record_id}.jpg",
            "audio_url": f"/media/{record_id}.wav",
        }
        if mutter_wav:
            public["mutter_audio"] = f"/media/{record_id}_mutter.wav"
        jpg, wav = self.store.write_media(record_id, image_jpeg, grudge_wav)
        dirty = [jpg, wav]
        if mutter_wav:
            path = self.store.write_blob(f"{record_id}_mutter.wav", mutter_wav)
            self._media[(f"{record_id}_mutter", "wav")] = path
            dirty.append(path)
        if og_jpeg:
            path = self.store.write_blob(f"{record_id}_og.jpg", og_jpeg)
            self._media[(f"{record_id}_og", "jpg")] = path
            dirty.append(path)
        self.store.append(public)
        self._records.append(public)
        self._media[(record_id, "jpg")] = jpg
        self._media[(record_id, "wav")] = wav
        self._mark_dirty(*dirty)
        return public

    # ----------------------------------------------------------------- dataset sync
    def _mark_dirty(self, *paths: Path) -> None:
        """Queue files for the next debounced batch; spawn the uploader task if
        none is in flight. No-op (and no growth) when sync is disabled."""
        if not self.sync.enabled:
            return
        self._dirty.update(paths)
        self._dirty.add(self.store.path)
        if self._sync_task is None or self._sync_task.done():
            self._sync_task = asyncio.create_task(self._upload_when_due())

    async def _upload_when_due(self) -> None:
        """Wait out the debounce (>=60s between commits), then push ONE batch of
        everything dirty. Anything marked dirty during the upload is picked up
        by a follow-up round; failures re-queue for the next publish."""
        while self._dirty:
            if self._last_upload is not None:
                delay = self.debounce - (self.clock() - self._last_upload)
                if delay > 0:
                    await asyncio.sleep(delay)
            batch = sorted(self._dirty)
            self._dirty.clear()
            ok = await asyncio.to_thread(self.sync.upload_batch, list(batch))
            self._last_upload = self.clock()
            if not ok:
                self._dirty.update(batch)  # retry rides on the next publish
                return

    async def drain(self) -> None:
        """Shutdown hygiene: stop the timer and push whatever is still dirty."""
        task = self._sync_task
        if task is not None and not task.done():
            task.cancel()
            await asyncio.gather(task, return_exceptions=True)
        if self.sync.enabled and self._dirty:
            batch = sorted(self._dirty)
            self._dirty.clear()
            await asyncio.to_thread(self.sync.upload_batch, list(batch))


# ------------------------------------------------------------------- pending vault
@dataclass
class PendingRecord:
    """One unpublished awakening, parked server-side between awaken and publish."""

    record: dict[str, Any]
    image_jpeg: bytes
    grudge_wav: bytes
    gate: dict[str, Any]
    created_at: float = 0.0
    # Best-effort sibling (contract): None when the GPU-window margin skipped
    # the mutter TTS or it failed — publish then simply persists no mutter.
    mutter_wav: Optional[bytes] = None


@dataclass
class PendingStore:
    """One-time tokens -> pending records. 15-minute TTL, single use (claim
    pops), oldest evicted beyond ``max_pending``. ``clock`` injectable for
    expiry tests. Single-threaded access (the asyncio request loop) by design.
    """

    ttl: float = PENDING_TTL_SECONDS
    max_pending: int = MAX_PENDING
    clock: Callable[[], float] = time.monotonic
    _items: dict[str, PendingRecord] = field(default_factory=dict)

    def put(
        self,
        record: dict[str, Any],
        image_jpeg: bytes,
        grudge_wav: bytes,
        gate: dict[str, Any],
        mutter_wav: Optional[bytes] = None,
    ) -> str:
        """Park a pending record; returns the fresh one-time token."""
        self._sweep()
        while len(self._items) >= self.max_pending:
            oldest = min(self._items, key=lambda k: self._items[k].created_at)
            del self._items[oldest]
        token = secrets.token_urlsafe(32)
        self._items[token] = PendingRecord(
            record=record,
            image_jpeg=image_jpeg,
            grudge_wav=grudge_wav,
            gate=gate,
            created_at=self.clock(),
            mutter_wav=mutter_wav,
        )
        return token

    def peek(self, token: str) -> Optional[PendingRecord]:
        """Read WITHOUT spending: lets the publish endpoint validate the token
        before its rate limiters record anything (security review #2c —
        unauthenticated junk-token floods must not inflate limiter state).
        Sweeps first, so an expired token is as invisible as an unknown one."""
        self._sweep()
        return self._items.get(token)

    def claim(self, token: str) -> Optional[PendingRecord]:
        """Pop the pending record for ``token`` — single use; None when the
        token is unknown, already used, or has faded past its TTL."""
        self._sweep()
        return self._items.pop(token, None)

    def __len__(self) -> int:
        self._sweep()
        return len(self._items)

    def _sweep(self) -> None:
        now = self.clock()
        stale = [t for t, p in self._items.items() if now - p.created_at >= self.ttl]
        for token in stale:
            del self._items[token]