src/ad_creative_env/renderer/cache.py

"""Two-tier background cache — disk only (no HTTP, no provider knowledge).

dev tier   = every network fill lands here (non-authoritative, re-fillable).
eval tier  = written ONLY by promote() (Task 8); replay/eval reads only here.

Storage is CONTENT-ADDRESSED with a single atomic commit point:
  * the image is an IMMUTABLE blob <key>.<byte_hash>.png — same bytes -> same
    name, so it is written once and NEVER overwritten in place;
  * the manifest <key>.json is the commit marker and REFERENCES that blob by its
    byte_hash. Publishing writes the blob first (durable), then atomically renames
    the manifest into place — that ONE manifest rename is the whole commit.

So a reader (which takes NO lock) always sees either the complete old entry or the
complete new entry — never a torn pair, even during an overwrite: a new publish
writes a NEW blob alongside the old one and flips the manifest atomically, and an
interrupted publish just leaves an orphan blob the still-old manifest never points
at (harmless — re-fillable in dev, and the frozen eval tier never overwrites).
has()/get() key off the manifest marker. Public put() is the DEV door only; the
eval trust boundary is crossed solely by promote().
"""

from __future__ import annotations

import contextlib
import fcntl
import hashlib
import io
import json
import logging
import os
import re
import uuid
from pathlib import Path

from PIL import Image

logger = logging.getLogger(__name__)


class CacheError(RuntimeError):
    """A cache integrity / policy violation (missing blob, freeze, identity, tier, key)."""


_TIERS = ("dev", "eval")

# A cache key must be a single flat path component. Production keys are always
# cache_key() output (64-hex), which matches; this is a defense-in-depth boundary
# check so a malformed/untrusted key can never traverse out of its tier dir (e.g.
# "../eval/x") and reach the eval tier through the dev-only put().
_KEY_RE = re.compile(r"\A[A-Za-z0-9_-]+\Z")


def _validate_key(key: str) -> str:
    if not isinstance(key, str) or not _KEY_RE.match(key):
        raise CacheError(f"invalid cache key {key!r}: must match [A-Za-z0-9_-]+")
    return key


def cache_key(fields: dict) -> str:
    """Stable SHA-256 over a canonicalized dict (sorted keys, compact)."""
    blob = json.dumps(fields, sort_keys=True, separators=(",", ":"))
    return hashlib.sha256(blob.encode("utf-8")).hexdigest()


def _byte_hash(data: bytes) -> str:
    return hashlib.sha256(data).hexdigest()


def _fsync(path: Path) -> None:
    """fsync a file or directory so a rename/write survives a crash (POSIX/Linux)."""
    fd = os.open(path, os.O_RDONLY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def _png_bytes(image: Image.Image) -> bytes:
    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="PNG")
    return buf.getvalue()


class Cache:
    def __init__(self, dev_dir, eval_dir, *, replay: bool = False):
        self.dev_dir = Path(dev_dir)
        self.eval_dir = Path(eval_dir)
        self.replay = replay
        self.dev_dir.mkdir(parents=True, exist_ok=True)
        self.eval_dir.mkdir(parents=True, exist_ok=True)
        # Aliased tier dirs (same path, or symlinks resolving to one) would let the
        # dev-only put() land in eval — collapsing the trust boundary. Reject it.
        if self.dev_dir.resolve() == self.eval_dir.resolve():
            raise CacheError(
                "dev and eval cache dirs must be distinct — "
                "aliasing them defeats the eval trust boundary"
            )

    def _dir(self, tier: str) -> Path:
        if tier not in _TIERS:
            raise CacheError(f"unknown cache tier {tier!r} (use 'dev' or 'eval')")
        return self.eval_dir if tier == "eval" else self.dev_dir

    def _manifest_path(self, key: str, tier: str) -> Path:
        return self._dir(tier) / f"{key}.json"  # the commit marker

    def _blob_path(self, key: str, tier: str, byte_hash: str) -> Path:
        return self._dir(tier) / f"{key}.{byte_hash}.png"  # immutable, content-addressed

    def has(self, key: str, tier: str) -> bool:
        _validate_key(key)
        return self._manifest_path(key, tier).exists()  # manifest = commit marker

    def put(self, key: str, image: Image.Image, manifest: dict) -> None:
        """Publish an entry to the DEV tier. Eval is written only by promote()."""
        _validate_key(key)
        self._publish(key, "dev", image, manifest)

    def _publish(self, key: str, tier: str, image: Image.Image, manifest: dict) -> None:
        """Publish image+manifest to `tier` with a single atomic commit.

        1. Materialize the image as an immutable content-addressed blob (idempotent:
           if the blob already exists its bytes are identical, so skip the write).
        2. Atomically rename the manifest (which references the blob by byte_hash)
           into place — this ONE rename commits the entry. A crash between the two
           steps leaves only an orphan blob, never a torn entry, and an overwrite
           adds a new blob rather than mutating the one the old manifest points at.

        Temp names use a uuid4 so uncoordinated writers (even two Cache instances in
        one process) never collide, and each temp lifecycle is wrapped in try/finally
        so a mid-write failure leaves no stray `.tmp` behind. Blob and manifest are
        fsync'd (contents + directory) so a committed manifest is never durable ahead
        of the bytes it references.
        """
        d = self._dir(tier)
        data = _png_bytes(image)
        byte_hash = _byte_hash(data)
        manifest = {**manifest, "byte_hash": byte_hash}
        blob_path = self._blob_path(key, tier, byte_hash)

        # Content-addressed: a blob already at this name is reused ONLY if its bytes
        # actually hash to that name. A corrupt blob squatting on the name (external
        # tampering / fs corruption) is rewritten with the good bytes we hold — never
        # trusted from the filename alone.
        needs_blob = True
        if blob_path.exists():
            needs_blob = _byte_hash(blob_path.read_bytes()) != byte_hash
        if needs_blob:
            tmp_blob = d / f"{blob_path.name}.{uuid.uuid4().hex}.tmp"
            try:
                tmp_blob.write_bytes(data)
                _fsync(tmp_blob)  # durable contents before any manifest points at it
                tmp_blob.replace(blob_path)
            finally:
                tmp_blob.unlink(missing_ok=True)  # only bites if we raised before replace
            _fsync(d)  # durable blob rename

        man_path = self._manifest_path(key, tier)
        tmp_man = d / f"{key}.json.{uuid.uuid4().hex}.tmp"
        try:
            tmp_man.write_text(json.dumps(manifest, indent=2))
            _fsync(tmp_man)
            tmp_man.replace(man_path)  # the single atomic commit point
        finally:
            tmp_man.unlink(missing_ok=True)
        _fsync(d)  # durable manifest commit
        logger.debug("cache write", extra={"tier": tier, "key": key[:12]})

    def get(
        self, key: str, tier: str, *, expect_slug: str | None = None
    ) -> tuple[Image.Image, dict] | None:
        """Return (image, manifest) or None on a clean miss (no commit marker).

        When the manifest marker IS present, hard-errors (CacheError) if it
        references a missing image blob (corruption), if the blob's bytes do not
        hash to the recorded byte_hash (tampering), or on a config-identity
        mismatch (expect_slug).
        """
        _validate_key(key)
        man_path = self._manifest_path(key, tier)
        if not man_path.exists():
            logger.debug("cache miss", extra={"tier": tier, "key": key[:12]})
            return None  # no commit marker = clean miss (orphan blob self-heals)
        try:
            manifest = json.loads(man_path.read_text())
        except (ValueError, OSError) as exc:  # corrupt/truncated manifest JSON
            # Inside an except -> exc_info=True captures the REAL cause (JSON/OS error).
            logger.error("cache guard: unreadable manifest",
                         extra={"tier": tier, "key": key[:12]}, exc_info=True)
            raise CacheError(f"unreadable manifest for {key} in {tier}") from exc
        byte_hash = manifest.get("byte_hash")
        blob_path = self._blob_path(key, tier, byte_hash or "")
        if not byte_hash or not blob_path.exists():
            # A DETECTED condition (no exception is being handled) -> structured ERROR,
            # NO exc_info (there is no live traceback; exc_info=True would log "NoneType").
            logger.error("cache guard: missing image blob",
                         extra={"tier": tier, "key": key[:12]})
            raise CacheError(
                f"manifest for {key} in {tier} references a missing image blob"
            )
        data = blob_path.read_bytes()
        if _byte_hash(data) != byte_hash:
            logger.error("cache guard: byte-hash mismatch",
                         extra={"tier": tier, "key": key[:12]})
            raise CacheError(f"byte-hash mismatch for {key} in {tier}")
        if expect_slug is not None and manifest.get("model_slug") != expect_slug:
            logger.error("cache guard: config-identity mismatch",
                         extra={"tier": tier, "key": key[:12],
                                "have": manifest.get("model_slug"), "want": expect_slug})
            raise CacheError(
                f"config-identity: entry slug {manifest.get('model_slug')!r} "
                f"!= configured {expect_slug!r}"
            )
        image = Image.open(io.BytesIO(data)).convert("RGB")
        logger.debug("cache hit", extra={"tier": tier, "key": key[:12]})
        return image, manifest

    def promote(self, key: str, *, is_eligible) -> None:
        """Freeze a dev entry into the eval tier — the ONLY eval-write path.

        The whole critical section (eligibility -> freeze check -> publish) runs
        under the EVAL-tier lock(key, "eval"), and eval is re-read while the lock
        is held (revalidate-while-locked). Locking the eval tier — not dev — is
        what serializes competing promoters even when they were constructed with
        different dev dirs but a shared eval dir, so two promoters can never both
        pass the freeze and publish different bytes: exactly one distinct content
        wins. The freeze compares the freshly re-encoded bytes (what _publish will
        write) against the frozen eval hash, so a lossless decode/re-encode never
        false-trips the freeze; a matching hash is an idempotent no-op that does
        NOT rewrite the already-frozen eval manifest.

        Precondition: the dev entry is expected to be already committed. promote
        holds only the eval lock, so it does NOT wait for an in-flight dev fill of
        the same key (callers fill-then-promote sequentially). If dev is absent it
        raises rather than blocking — a safe failure, never corruption.
        """
        _validate_key(key)
        with self.lock(key, "eval"):
            entry = self.get(key, "dev")
            if entry is None:
                raise CacheError(f"cannot promote {key}: not in dev cache")
            image, manifest = entry
            slug = manifest.get("model_slug")
            if not is_eligible(slug):
                raise CacheError(f"cannot promote {key}: model {slug!r} not promotion-eligible")

            existing = self.get(key, "eval")  # revalidate under the held lock
            if existing is not None:
                new_hash = _byte_hash(_png_bytes(image))  # bytes _publish would write
                if existing[1].get("byte_hash") != new_hash:
                    raise CacheError(f"eval freeze: {key} differs from the frozen eval entry")
                logger.debug("promote no-op: already frozen", extra={"key": key[:12]})
                return  # identical bytes: frozen entry stands — never rewrite its manifest
            self._publish(key, "eval", image, manifest)
            logger.info("promoted to eval", extra={"key": key[:12], "model_slug": slug})

    @contextlib.contextmanager
    def lock(self, key: str, tier: str = "dev"):
        """Exclusive per-key lock (fcntl.flock) scoped to a tier's directory.

        Cross-process mutual exclusion so parallel env workers don't make duplicate
        paid calls or race on the same key. The lock lives in the directory of the
        tier it guards — dev fills lock `dev` (default), promote locks `eval` — so
        instances sharing that tier's dir serialize even if their other dir differs.
        POSIX/Linux.
        """
        _validate_key(key)
        lock_path = self._dir(tier) / f"{key}.lock"
        handle = open(lock_path, "w")
        try:
            fcntl.flock(handle, fcntl.LOCK_EX)
            yield
        finally:
            fcntl.flock(handle, fcntl.LOCK_UN)
            handle.close()