"""
pipeline_services/common/pipeline.py

Shared definitions for all pipeline stage services.

Stream topology (flat namespace, separate Redis streams):
    stage.raw           → tokenizer_service      → stage.tokenized
    stage.tokenized     → forward_pass_service   → stage.hidden
    stage.hidden        → chunker_service        → stage.chunked
    stage.chunked       → extractor_service      → stage.embedded
    stage.embedded      → tagger_service         → stage.tagged
    stage.tagged        → storage_service        → stage.done
    stage.done          → compressor_service     → (compresses stage files in Garage)

Message format (always lightweight — payload lives in Garage S3):
    {
        "doc_id":       str,          # UUID for this document
        "stage":        str,          # current stage name
        "storage_key":  str,          # Garage S3 key for the payload
        "context_uuid": str,          # UUID-zero or user profile UUID
        "metadata":     JSON str,     # title, author, year, tags etc.
    }

Storage layout in Garage S3 bucket "data":
    pipeline/{context_uuid}/{doc_id}/meta.json           ← always uncompressed, permanent
    pipeline/{context_uuid}/{doc_id}/stage.raw.json
    pipeline/{context_uuid}/{doc_id}/stage.tokenized.json
    pipeline/{context_uuid}/{doc_id}/stage.hidden.json
    pipeline/{context_uuid}/{doc_id}/stage.chunked.json
    pipeline/{context_uuid}/{doc_id}/stage.embedded.json
    pipeline/{context_uuid}/{doc_id}/stage.tagged.json
    → after compressor_service runs:
    pipeline/{context_uuid}/{doc_id}/stage.hidden.json.bz2
    (etc. for all completed stages)

Checkpoint behaviour:
    Before running any stage, check if the output storage_key already exists.
    If it does — skip this stage and publish the next stage message directly.
    This means a failed job can be resubmitted and resumes from the last
    completed stage. The expensive Qwen forward pass is never repeated.
"""
from __future__ import annotations

import bz2
import json
import logging
import os
import shutil
import subprocess
import tempfile
import time
import uuid
from typing import Any, BinaryIO, Optional

import boto3
import redis as _redis_module
from botocore.config import Config

logger = logging.getLogger(__name__)

# ── Stream names ───────────────────────────────────────────────────────────────

STREAMS = {
    "raw":       "stage.raw",
    "tokenized": "stage.tokenized",
    "hidden":    "stage.hidden",
    "chunked":   "stage.chunked",
    "embedded":  "stage.embedded",
    "tagged":    "stage.tagged",
    "done":      "stage.done",
    "compress":  "stage.compress",
}

CONSUMER_GROUPS = {
    "stage.raw":       "tokenizer_workers",
    "stage.tokenized": "forward_pass_workers",
    "stage.hidden":    "chunker_workers",
    "stage.chunked":   "extractor_workers",
    "stage.embedded":  "tagger_workers",
    "stage.tagged":    "storage_workers",
    "stage.done":      "compressor_workers",
    "stage.compress":  "compressor_workers",
}

GARAGE_BUCKET = "data"
PIPELINE_PREFIX = "pipeline"

# ── Environment ────────────────────────────────────────────────────────────────

REDIS_URL        = os.getenv("REDIS_URL",           "redis://10.121.1.60:6379/0")
# NB: :3900 is Garage's S3 API. :3901 is the cluster RPC port — pointing the S3
# client there silently fails (RPC closes the HTTP connection, surfacing as a
# botocore ConnectionClosedError). A remote pool worker with no GARAGE_URL env
# faceplanted on the old :3901 default before this was corrected.
# ── S3 credentials: config-service FIRST, env fallback (operator 2026-07-17) ─────
# Retiring the /etc/andrea dot-envs as the key SOURCE: a Garage-key rotation should
# propagate from ONE place (the central config-service) instead of leaving stale keys
# baked into worker envs/images (which bit us — a rotated gitlab key never reached the
# forward_pass workers). config_client is STDLIB-only + FAIL-SOFT (service down/empty →
# {}), so the env fallback preserves TODAY's exact behaviour when the config-service says
# nothing. Resolved ONCE at import (the clients below are lazy); a worker restart
# re-resolves — so flipping to the config value is a RESTART, never a re-bake. The
# per-box row overlays the global '*' (savvy boxes can carry their own key later).
def _pipeline_s3_cfg() -> dict:
    try:
        try:
            from common import config_client
        except ImportError:
            import config_client               # common/ directly on sys.path (baked worker images)
        box = os.getenv("WORKER_BOX", "").strip() or None
        return config_client.fetch_config("pipeline-s3", box) or {}
    except Exception:                          # import must NEVER depend on the config-service
        return {}


_S3_CFG = _pipeline_s3_cfg()


def _s3cfg(key: str, default: str = "") -> str:
    """config-service value (present + non-empty) wins; else env; else default."""
    v = _S3_CFG.get(key)
    return v if v else os.getenv(key, default)


GARAGE_URL       = _s3cfg("GARAGE_URL",          "http://10.121.1.1:3900")
GARAGE_KEY_ID    = _s3cfg("GARAGE_GITLAB_KEY_ID", "")
GARAGE_SECRET    = _s3cfg("GARAGE_GITLAB_SECRET", "")
GARAGE_REGION    = _s3cfg("GARAGE_REGION",        "garage")

# ── Global-fallback store (savvy cutover) ───────────────────────────────────────
# When a worker's GARAGE_URL points at a box-local SAVVY cluster, an input atom may
# not yet be resident there — it lives in the GLOBAL Garage (the durable home, where
# the doc's GUID always exists). The read path therefore resolves SAVVY-first → on a
# miss read GLOBAL → COPY the atom INTO savvy ("start a sequence" box-local) → return,
# and NEVER errors on a savvy miss (operator 2026-06-21). GLOBAL_* default to GARAGE_*
# so a non-cutover box (global == savvy) makes the whole helper a pure no-op.
GLOBAL_S3_URL    = _s3cfg("GLOBAL_S3_URL",        GARAGE_URL)
GLOBAL_KEY_ID    = _s3cfg("GLOBAL_KEY_ID",        GARAGE_KEY_ID)
GLOBAL_SECRET    = _s3cfg("GLOBAL_SECRET",        GARAGE_SECRET)

# ── PRIMARY-store backend selector (S3/Garage vs box-local FOLDERS) ──────────────
# The PRIMARY (savvy) store can be a box-local FILESYSTEM instead of an S3/Garage
# cluster, switched SILENTLY by config (operator 2026-06-23): "no savvy uri provided
# → the cross-tb S3 savvy code goes dormant; we're on the folders; remove the folder
# path → that code goes dormant." When SAVVY_LOCAL_DIR is non-empty the PRIMARY store
# is the filesystem rooted there (key K → os.path.join(SAVVY_LOCAL_DIR, K)); the GLOBAL
# durable store stays S3 (GLOBAL_S3_URL). When SAVVY_LOCAL_DIR is unset, every S3
# codepath below is BYTE-IDENTICAL to before this backend existed.
#
# On-disk layout ASSUMPTION (flag for the operator): the folder MIRRORS the S3 key
# layout — pipeline/<context_uuid>/<doc_id>/stage.<stage>.json under SAVVY_LOCAL_DIR.
# A different on-disk structure would need a key→path mapping override here.
SAVVY_LOCAL_DIR  = os.getenv("SAVVY_LOCAL_DIR",      "").strip()

# True when the PRIMARY store differs from the GLOBAL durable store — i.e. the box is
# flipped to savvy. Two ways to flip: (a) GARAGE_URL points at a box-local savvy S3
# cluster distinct from GLOBAL_S3_URL, or (b) SAVVY_LOCAL_DIR is set (folder mode — the
# primary is the filesystem, the global is still S3, so savvy-first + copy-in apply and
# copy-in pulls GLOBAL S3 → local file).
SAVVY_ENABLED    = bool(SAVVY_LOCAL_DIR) or (GLOBAL_S3_URL != GARAGE_URL)

# Garage v2.2.0 + botocore ≥1.36 default CRC32 streaming breaks reads/copies of large
# atoms (FlexibleChecksumError on a ~510 MB .npy, verified live 2026-06-20). Pin
# checksums to when_required at the ENV level BEFORE any boto3 client is built
# (pipeline is imported at worker top; clients are lazy) — matching the savvy services.
os.environ.setdefault("AWS_RESPONSE_CHECKSUM_VALIDATION", "when_required")
os.environ.setdefault("AWS_REQUEST_CHECKSUM_CALCULATION", "when_required")

# ── compressed-atom read-path fallback (mirror of the compressor's .bz2 write) ───
# download_atom transparently falls back to <key>.bz2 on a raw-key miss, decompresses
# it STREAMING (chunk-by-chunk — never buffers the multi-GB tensor in RAM) into the
# caller's fileobj. ATOM_BZ2_CHUNK is the streaming-copy granularity (compressed-side
# read for the python path, decompressed-side write for the pbzip2 path).
ATOM_BZ2_CHUNK    = int(os.getenv("ATOM_BZ2_CHUNK", str(8 * 1024 ** 2)))   # 8 MiB
# pbzip2 -d is preferred (parallel) when present; python bz2 is the always-available
# fallback. Empty PBZIP2_BIN disables the subprocess path (force pure-python decompress).
PBZIP2_BIN        = os.getenv("PBZIP2_BIN", "pbzip2").strip()
# Where the compressed .bz2 is buffered before decompression. Mirrors the compressor's
# tmpfs-gate IN REVERSE: when the caller's fileobj lands on rotational md1 the .bz2
# scratch goes to tmpfs (sequential RAM read for the decompressor) when it FITS; else
# /tmp. Empty → use the system default temp dir (tempfile picks /tmp).
ATOM_BZ2_TMP_DIR  = os.getenv("ATOM_BZ2_TMP_DIR", "").strip()
ATOM_BZ2_TMPFS_DIR = os.getenv("ATOM_BZ2_TMPFS_DIR", "/dev/shm").strip()
# Stage the compressed .bz2 through tmpfs only when it is at least this big AND fits.
ATOM_BZ2_TMPFS_MIN_BYTES = int(os.getenv("ATOM_BZ2_TMPFS_MIN_BYTES", str(256 * 1024 ** 2)))  # 256 MiB
ATOM_BZ2_TMPFS_SAFETY    = float(os.getenv("ATOM_BZ2_TMPFS_SAFETY", "1.2"))


# ── Redis connection ───────────────────────────────────────────────────────────

_redis_client: Optional[_redis_module.Redis] = None


def get_redis() -> _redis_module.Redis:
    global _redis_client
    if _redis_client is None:
        _redis_client = _redis_module.from_url(
            REDIS_URL,
            decode_responses=True,
            socket_connect_timeout=5,
            socket_timeout=60,
            retry_on_timeout=True,
        )
        _redis_client.ping()
        logger.info("Redis connected: %s", REDIS_URL)
    return _redis_client


def ensure_streams(r: Optional[_redis_module.Redis] = None) -> None:
    """Create all streams and consumer groups if they don't exist. Idempotent."""
    r = r or get_redis()
    for stream, group in CONSUMER_GROUPS.items():
        try:
            r.xgroup_create(stream, group, id="0", mkstream=True)
            logger.debug("Created group %r on %r", group, stream)
        except _redis_module.exceptions.ResponseError as e:
            if "BUSYGROUP" not in str(e):
                raise


# ── Garage S3 client ───────────────────────────────────────────────────────────

_s3_client = None
_global_s3_client = None


def _garage_cfg():
    """boto3 Config for Garage: s3v4 + checksum=when_required (the CRC32/FlexibleChecksum
    trap on big atoms) + standard retries. Degrades gracefully on botocore too old for the
    checksum kwargs."""
    base = dict(signature_version="s3v4", retries={"max_attempts": 3, "mode": "standard"})
    try:
        return Config(response_checksum_validation="when_required",
                      request_checksum_calculation="when_required", **base)
    except TypeError:
        return Config(**base)


def get_s3():
    """The worker's PRIMARY store client (GARAGE_URL — box-local savvy when flipped, else global)."""
    global _s3_client
    if _s3_client is None:
        _s3_client = boto3.client(
            "s3",
            endpoint_url=GARAGE_URL,
            aws_access_key_id=GARAGE_KEY_ID,
            aws_secret_access_key=GARAGE_SECRET,
            region_name=GARAGE_REGION,
            config=_garage_cfg(),
        )
    return _s3_client


def get_global_s3():
    """The GLOBAL (durable home) store client — the savvy-miss read SOURCE. When SAVVY is
    not enabled (global == savvy) this is just get_s3()."""
    global _global_s3_client
    if not SAVVY_ENABLED:
        return get_s3()
    if _global_s3_client is None:
        _global_s3_client = boto3.client(
            "s3",
            endpoint_url=GLOBAL_S3_URL,
            aws_access_key_id=GLOBAL_KEY_ID,
            aws_secret_access_key=GLOBAL_SECRET,
            region_name=GARAGE_REGION,
            config=_garage_cfg(),
        )
    return _global_s3_client


def _is_not_found(exc) -> bool:
    """True if a botocore error is a 404 / NoSuchKey / NotFound miss."""
    if isinstance(exc, (FileNotFoundError, _LocalNotFound)):   # folder-backend miss
        return True
    code = ""
    resp = getattr(exc, "response", None)
    if isinstance(resp, dict):
        code = resp.get("Error", {}).get("Code", "")
    return code in ("NoSuchKey", "NotFound", "404")


# ── PRIMARY-store interface (S3/Garage  ⟷  box-local FOLDERS) ────────────────────
# A minimal store contract the PRIMARY-store call sites speak. Two implementations:
#   _S3Store     — wraps the existing boto3 client; behaviour is BYTE-IDENTICAL to the
#                  pre-existing direct get_s3() calls (it IS get_s3()).
#   _FolderStore — a box-local filesystem rooted at SAVVY_LOCAL_DIR, key K mapped to
#                  os.path.join(root, K). Streams large atoms (write to <path>.part +
#                  os.replace for atomicity); list = os.walk; head = exists + st_size.
# The PRIMARY impl is selected ONCE at import (folder if SAVVY_LOCAL_DIR else S3); the
# GLOBAL store is ALWAYS S3 (the durable home / copy-in source).

class _LocalNotFound(Exception):
    """Folder-backend equivalent of a NoSuchKey miss (mapped by _is_not_found)."""


class _S3Store:
    """PRIMARY/GLOBAL store backed by a boto3 S3 client. Thin pass-through — every
    method maps 1:1 onto the boto3 call the old direct code made, so the S3 path stays
    byte-identical. `client` is resolved lazily (the boto3 client is built on first use,
    preserving the original lazy-client semantics)."""

    def __init__(self, client_fn, bucket: str = GARAGE_BUCKET):
        self._client_fn = client_fn
        self.bucket = bucket

    @property
    def client(self):
        return self._client_fn()

    def put_bytes(self, key: str, data: bytes, content_type: str = "application/octet-stream",
                  metadata: Optional[dict] = None) -> None:
        extra = {}
        if metadata:
            extra["Metadata"] = metadata
        self.client.put_object(Bucket=self.bucket, Key=key, Body=data,
                               ContentType=content_type, **extra)

    def put_stream(self, key: str, fileobj: BinaryIO, content_type: str = "application/octet-stream",
                   metadata: Optional[dict] = None) -> None:
        extra = {"ContentType": content_type}
        if metadata:
            extra["Metadata"] = metadata
        self.client.upload_fileobj(fileobj, self.bucket, key, ExtraArgs=extra)

    def get_bytes(self, key: str) -> bytes:
        return self.client.get_object(Bucket=self.bucket, Key=key)["Body"].read()

    def get_stream(self, key: str, fileobj: BinaryIO) -> None:
        self.client.download_fileobj(self.bucket, key, fileobj)

    def head(self, key: str) -> Optional[int]:
        """Return ContentLength if present, else raise a not-found error."""
        return self.client.head_object(Bucket=self.bucket, Key=key).get("ContentLength")

    def exists(self, key: str) -> bool:
        try:
            self.client.head_object(Bucket=self.bucket, Key=key)
            return True
        except Exception:
            return False

    def list(self, prefix: str):
        paginator = self.client.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=self.bucket, Prefix=prefix):
            for o in page.get("Contents", []):
                yield o["Key"]

    def delete(self, key: str) -> None:
        self.client.delete_object(Bucket=self.bucket, Key=key)


class _FolderStore:
    """PRIMARY store backed by a box-local directory tree rooted at `root`. Key K maps to
    os.path.join(root, K) — the folder MIRRORS the S3 key layout (pipeline/<ctx>/<doc>/…).

    • put_*       — stream/write to <path>.part then os.replace (atomic publish); parents
                    created mkdir -p; bytes stored VERBATIM (the graduation sha256 is over
                    the bytes, so the on-disk file IS the verifiable atom).
    • metadata    — when non-empty, written to a sidecar <path>.meta.json (ContentType +
                    user Metadata). Load-bearing only where graduation reads the producer
                    content-sha256 sidecar; the atom bytes themselves carry no framing.
    • head        — exists + st_size  (mirrors S3 head_object ContentLength).
    • list        — os.walk under the prefix dir, yielding S3-style keys (root-relative,
                    forward-slashed), incl. a file whose key == prefix exactly.
    • delete      — unlink atom + its sidecar (best-effort)."""

    META_SUFFIX = ".meta.json"

    def __init__(self, root: str):
        self.root = os.path.abspath(root)

    def _path(self, key: str) -> str:
        # Keys are forward-slashed, root-relative, never absolute or '..'-escaping.
        return os.path.join(self.root, *key.split("/"))

    def _write_sidecar(self, path: str, content_type: str, metadata: Optional[dict]) -> None:
        side = {"ContentType": content_type}
        if metadata:
            side["Metadata"] = metadata
        tmp = path + self.META_SUFFIX + ".part"
        with open(tmp, "w", encoding="utf-8") as fh:
            json.dump(side, fh, ensure_ascii=False)
        os.replace(tmp, path + self.META_SUFFIX)

    def put_bytes(self, key: str, data: bytes, content_type: str = "application/octet-stream",
                  metadata: Optional[dict] = None) -> None:
        path = self._path(key)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        tmp = path + ".part"
        with open(tmp, "wb") as fh:
            fh.write(data)
        os.replace(tmp, path)                        # atomic publish
        if metadata:
            self._write_sidecar(path, content_type, metadata)

    def put_stream(self, key: str, fileobj: BinaryIO, content_type: str = "application/octet-stream",
                   metadata: Optional[dict] = None) -> None:
        path = self._path(key)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        tmp = path + ".part"
        with open(tmp, "wb") as fh:
            shutil.copyfileobj(fileobj, fh)          # streamed — never buffers the atom in RAM
        os.replace(tmp, path)
        if metadata:
            self._write_sidecar(path, content_type, metadata)

    def get_bytes(self, key: str) -> bytes:
        try:
            with open(self._path(key), "rb") as fh:
                return fh.read()
        except FileNotFoundError as e:
            raise _LocalNotFound(key) from e

    def get_stream(self, key: str, fileobj: BinaryIO) -> None:
        try:
            with open(self._path(key), "rb") as fh:
                shutil.copyfileobj(fh, fileobj)
        except FileNotFoundError as e:
            raise _LocalNotFound(key) from e

    def head(self, key: str) -> int:
        try:
            return os.stat(self._path(key)).st_size
        except FileNotFoundError as e:
            raise _LocalNotFound(key) from e

    def exists(self, key: str) -> bool:
        return os.path.isfile(self._path(key))

    def list(self, prefix: str):
        base = self._path(prefix)
        if os.path.isfile(base):                     # prefix names an exact object
            yield prefix
            return
        if not os.path.isdir(base):
            return
        for dirpath, _dirs, files in os.walk(base):
            for name in files:
                if name.endswith(self.META_SUFFIX) or name.endswith(".part"):
                    continue                         # sidecars / partials are not atoms
                full = os.path.join(dirpath, name)
                rel = os.path.relpath(full, self.root)
                yield rel.replace(os.sep, "/")

    def delete(self, key: str) -> None:
        path = self._path(key)
        for p in (path, path + self.META_SUFFIX):
            try:
                os.remove(p)
            except FileNotFoundError:
                pass


# PRIMARY store: folder when SAVVY_LOCAL_DIR is set, else the existing S3 client
# (byte-identical). GLOBAL store is ALWAYS S3 (durable home / copy-in source).
def _build_primary_store():
    if SAVVY_LOCAL_DIR:
        logger.info("PRIMARY store: FOLDER backend rooted at %s (savvy-on-folders)", SAVVY_LOCAL_DIR)
        return _FolderStore(SAVVY_LOCAL_DIR)
    return _S3Store(get_s3)


_primary_store = None
_global_store = None


def get_primary_store():
    """The worker's PRIMARY store (folder if SAVVY_LOCAL_DIR set, else S3 — unchanged)."""
    global _primary_store
    if _primary_store is None:
        _primary_store = _build_primary_store()
    return _primary_store


def get_global_store():
    """The GLOBAL (durable home) store — ALWAYS S3; the savvy-miss read SOURCE / copy-in
    source. When SAVVY is not enabled this wraps the SAME client as the primary (no-op)."""
    global _global_store
    if _global_store is None:
        _global_store = _S3Store(get_global_s3)
    return _global_store


def _copy_in_object(key: str) -> bool:
    """Idempotently copy one atom GLOBAL → box-local SAVVY ("start a sequence" box-local).
    STREAMS (never buffers the multi-GB atom in RAM); preserves ContentType + user Metadata
    so the provisioner's content-sha256 graduation verify still matches (S3: object metadata;
    folder: a sidecar .meta.json next to the atom). Returns True if the atom is now in savvy
    (copied or already present), False if it is absent in GLOBAL too. In folder mode the
    savvy side is the local filesystem and the global side is still S3, so copy-in pulls
    GLOBAL S3 → local file."""
    if not SAVVY_ENABLED:
        return False
    primary = get_primary_store()
    if primary.exists(key):                         # already in savvy? short-circuit (idempotent)
        return True
    try:                                            # pull from global S3 …
        src = get_global_s3().get_object(Bucket=GARAGE_BUCKET, Key=key)
    except Exception as e:
        if _is_not_found(e):
            return False                            # absent in global too — genuine miss
        raise
    content_type = src.get("ContentType", "application/octet-stream")
    metadata = src.get("Metadata") or None
    primary.put_stream(key, src["Body"], content_type=content_type, metadata=metadata)  # … push to savvy
    logger.info("copy-in global→savvy: %s", key)
    return True


def storage_key(context_uuid: str, doc_id: str, stage: str) -> str:
    return f"{PIPELINE_PREFIX}/{context_uuid}/{doc_id}/stage.{stage}.json"


def meta_key(context_uuid: str, doc_id: str) -> str:
    return f"{PIPELINE_PREFIX}/{context_uuid}/{doc_id}/meta.json"


def stage_exists(context_uuid: str, doc_id: str, stage: str) -> bool:
    """Check if a stage output already exists — checkpoint test. SAVVY-first then GLOBAL
    (a checkpoint resident only in global must read as present so the stage is not
    re-run). NO copy-in: a presence probe consumes no atom, so the pull fires only on
    the real read."""
    key = storage_key(context_uuid, doc_id, stage)
    stores = (get_primary_store(), get_global_store()) if SAVVY_ENABLED else (get_primary_store(),)
    for k in (key, key + ".bz2"):                   # also check the compressed variant
        for store in stores:
            if store.exists(k):
                return True
    return False


def write_stage(context_uuid: str, doc_id: str, stage: str, payload: Any) -> str:
    """Serialize payload as JSON and write to Garage. Returns the storage key."""
    key = storage_key(context_uuid, doc_id, stage)
    data = json.dumps(payload, ensure_ascii=False).encode("utf-8")
    get_primary_store().put_bytes(key, data, content_type="application/json")
    logger.debug("Wrote stage %s → %s (%d bytes)", stage, key, len(data))
    return key


def read_stage(storage_key_path: str) -> Any:
    """Read and deserialize a stage (JSON-envelope) payload. SAVVY-first → on a savvy miss
    read GLOBAL + best-effort copy-in → return. Never errors on a savvy miss when the atom
    is in global (the global read is the source of truth; copy-in starts the box-local
    sequence for the next reader)."""
    try:
        return json.loads(get_primary_store().get_bytes(storage_key_path).decode("utf-8"))
    except Exception as e:
        if not (SAVVY_ENABLED and _is_not_found(e)):
            raise
    resp = get_global_s3().get_object(Bucket=GARAGE_BUCKET, Key=storage_key_path)
    payload = json.loads(resp["Body"].read().decode("utf-8"))
    try:
        _copy_in_object(storage_key_path)
    except Exception as ce:                          # copy-in is best-effort; the global read already succeeded
        logger.warning("copy-in failed for %s (proceeding on global read): %s", storage_key_path, ce)
    return payload


def head_atom(key: str) -> bool:
    """Presence test for an arbitrary atom key. SAVVY-first then GLOBAL when flipped. No copy-in."""
    stores = (get_primary_store(), get_global_store()) if SAVVY_ENABLED else (get_primary_store(),)
    for store in stores:
        if store.exists(key):
            return True
    return False


def head_atom_size(key: str):
    """Byte size of an atom (SAVVY-first then GLOBAL when flipped), or None if absent.
    A presence+size probe — no copy-in."""
    stores = (get_primary_store(), get_global_store()) if SAVVY_ENABLED else (get_primary_store(),)
    for store in stores:
        try:
            return store.head(key)
        except Exception:
            continue
    return None


def list_atoms(prefix: str):
    """Enumerate atom keys under `prefix`, UNIONED across SAVVY-first then GLOBAL
    (savvy-resident-only AND global-only atoms both appear; a key present in both
    is yielded once). No copy-in — a listing consumes no atom. When SAVVY is not
    enabled this lists the single primary store. Order is unspecified; callers that
    care must sort. Used by the compressor to find a doc's per-window FP resume
    checkpoints (stage.hidden.json.winNNN.pt), of which there can be MANY per doc."""
    stores = (get_primary_store(), get_global_store()) if SAVVY_ENABLED else (get_primary_store(),)
    seen = set()
    for store in stores:
        try:
            for key in store.list(prefix):
                if key not in seen:
                    seen.add(key)
                    yield key
        except Exception as e:
            logger.warning("list_atoms: store list failed for prefix %s: %s", prefix, e)
            continue


def _stream_atom_raw(key: str, fileobj) -> None:
    """Stream the LITERAL atom `key` into fileobj via the savvy-first → copy-in → global
    machinery. SAVVY-first → on a savvy not-found copy-in from global → stream from savvy
    (now resident); if absent in global too, the final global stream raises NoSuchKey
    (genuine miss). Never buffers the atom in RAM. This is the EXACT pre-fallback body of
    download_atom — the raw fast path is byte-identical to before."""
    try:
        get_primary_store().get_stream(key, fileobj)
        return
    except Exception as e:
        if not (SAVVY_ENABLED and _is_not_found(e)):
            raise
    try:                                             # discard any partial bytes from the failed savvy attempt
        fileobj.seek(0); fileobj.truncate(0)
    except Exception:
        pass
    if _copy_in_object(key):                         # now resident in savvy → stream box-local
        get_primary_store().get_stream(key, fileobj)
    else:                                            # absent in global too → stream global (raises the genuine miss)
        get_global_s3().download_fileobj(GARAGE_BUCKET, key, fileobj)


def _file_is_rotational(path: str) -> bool:
    """True iff the block device backing `path` is ROTATIONAL (HDD / md1). Mount → device
    → base disk → /sys/block/<base>/queue/rotational (1=HDD). Anything unresolvable →
    False (treat as non-rotational; the SAFE default = no tmpfs staging). Self-contained
    mirror of the compressor's _is_rotational so common/ stays dependency-free."""
    try:
        target = os.path.abspath(path)
        best_mp, best_dev = "", None
        with open("/proc/mounts", "r") as fh:
            for line in fh:
                parts = line.split()
                if len(parts) < 2 or not parts[0].startswith("/dev/"):
                    continue
                dev, mp = parts[0], parts[1].replace("\\040", " ")
                if target == mp or target.startswith(mp.rstrip("/") + "/") or mp == "/":
                    if len(mp) >= len(best_mp):
                        best_mp, best_dev = mp, dev
        if not best_dev:
            return False
        name = os.path.basename(os.path.realpath(best_dev))

        def _resolve(n, depth=0):
            if not n or depth > 8:
                return None
            if os.path.exists(f"/sys/block/{n}/queue/rotational"):
                return n
            slaves_dir = f"/sys/class/block/{n}/slaves"
            if os.path.isdir(slaves_dir):
                slaves = [s for s in os.listdir(slaves_dir) if not s.startswith(".")]
                return _resolve(slaves[0], depth + 1) if len(slaves) == 1 else None
            part = f"/sys/class/block/{n}"
            if os.path.exists(f"{part}/partition"):
                parent = os.path.basename(os.path.dirname(os.path.realpath(part)))
                if parent and os.path.exists(f"/sys/block/{parent}/queue/rotational"):
                    return parent
            return None

        base = _resolve(name)
        if not base:
            return False
        with open(f"/sys/block/{base}/queue/rotational", "r") as fh:
            return fh.read().strip() == "1"
    except Exception:
        return False


def _bz2_scratch_dir(out_fileobj, compressed_size: Optional[int]) -> Optional[str]:
    """Where to buffer the compressed .bz2 before decompression. Mirrors the compressor's
    tmpfs-gate IN REVERSE: when the caller's output lands on a ROTATIONAL disk (md1), stage
    the .bz2 in tmpfs so the decompressor reads it sequentially from RAM — but ONLY when it
    is ≥ MIN and FITS (size*SAFETY ≤ min(tmpfs free, MemAvailable)). Returns the chosen dir,
    or None to let tempfile pick the default. ATOM_BZ2_TMP_DIR overrides everything."""
    if ATOM_BZ2_TMP_DIR:
        return ATOM_BZ2_TMP_DIR
    # Only worth staging when the decompressed output lands on a slow rotational disk.
    out_dir = None
    try:
        name = getattr(out_fileobj, "name", None)
        if isinstance(name, str) and name:
            out_dir = os.path.dirname(os.path.abspath(name)) or "/tmp"
    except Exception:
        out_dir = None
    if not out_dir or not _file_is_rotational(out_dir):
        return None
    if not ATOM_BZ2_TMPFS_DIR or not os.path.isdir(ATOM_BZ2_TMPFS_DIR):
        return None
    if compressed_size is None or compressed_size < ATOM_BZ2_TMPFS_MIN_BYTES:
        return None
    try:                                             # FITS check — never risk OOM on the tmpfs
        st = os.statvfs(ATOM_BZ2_TMPFS_DIR)
        shm_free = st.f_bavail * st.f_frsize
        mem_avail = 0
        with open("/proc/meminfo", "r") as fh:
            for line in fh:
                if line.startswith("MemAvailable:"):
                    mem_avail = int(line.split()[1]) * 1024
                    break
        if compressed_size * ATOM_BZ2_TMPFS_SAFETY > min(shm_free, mem_avail):
            return None
    except Exception:
        return None
    return ATOM_BZ2_TMPFS_DIR


def _decompress_bz2_stream(in_path: str, fileobj) -> None:
    """Decompress the bz2 file at `in_path` into `fileobj`, STREAMING chunk-by-chunk —
    NEVER buffering the multi-GB decompressed tensor in RAM. pbzip2 -d (parallel) when the
    binary is configured + present; otherwise python's bz2 streaming reader. Both write the
    decompressed bytes straight into fileobj in ATOM_BZ2_CHUNK slices."""
    if PBZIP2_BIN and shutil.which(PBZIP2_BIN):
        # pbzip2 -d -c <in_path>  → stdout, copied chunk-by-chunk into fileobj.
        proc = subprocess.Popen([PBZIP2_BIN, "-d", "-c", in_path],
                                stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        try:
            while True:
                chunk = proc.stdout.read(ATOM_BZ2_CHUNK)
                if not chunk:
                    break
                fileobj.write(chunk)
        finally:
            try:
                proc.stdout.close()
            except Exception:
                pass
        _err = proc.stderr.read() if proc.stderr else b""
        try:
            proc.stderr.close()
        except Exception:
            pass
        rc = proc.wait()
        if rc != 0:
            raise RuntimeError(f"pbzip2 -d failed (rc={rc}) on {in_path}: "
                               f"{_err.decode('utf-8', 'replace').strip()}")
        return
    # Pure-python streaming fallback — bz2.BZ2File is a multi-stream-safe incremental reader.
    with bz2.BZ2File(in_path, "rb") as bf:
        while True:
            chunk = bf.read(ATOM_BZ2_CHUNK)
            if not chunk:
                break
            fileobj.write(chunk)


def download_atom(key: str, fileobj) -> None:
    """Stream a large atom (.npy / .win{NNN}.pt / anything) into fileobj. SAVVY-first → on a
    savvy miss copy-in from global → stream from savvy (now resident). Never buffers the
    multi-GB atom in RAM.

    TRANSPARENT .bz2 FALLBACK (read-path mirror of the compressor's write): the LITERAL key
    is tried FIRST exactly as before (raw fast path 100% unchanged for existing raw atoms).
    Only on a NOT-FOUND miss of the raw key do we retry <key>+'.bz2' through the SAME
    savvy-first → copy-in → global machinery into a temp file, then DECOMPRESS it STREAMING
    (chunk-by-chunk — never buffers the decompressed tensor in RAM) into fileobj. Generic by
    key suffix, so the chunker/extractor call sites (download_atom(hidden_key, f) → np.load)
    inherit the fallback with ZERO changes.

    If neither the raw key NOR <key>.bz2 exists in savvy/global, the original NOT-FOUND error
    propagates (genuine miss → task fails + re-enqueues, per the durable-task contract)."""
    try:
        _stream_atom_raw(key, fileobj)
        return
    except Exception as e:
        if not _is_not_found(e):                     # a real error (perms/network/etc.) — surface it
            raise
        raw_miss = e                                 # remember the genuine raw miss to re-raise

    # Raw key absent everywhere → try the compressed sibling through the SAME machinery.
    bz2_key = key + ".bz2"
    try:                                             # discard any partial bytes from the raw attempt
        fileobj.seek(0); fileobj.truncate(0)
    except Exception:
        pass

    compressed_size = head_atom_size(bz2_key)
    scratch_dir = _bz2_scratch_dir(fileobj, compressed_size)
    tmp_fd, tmp_path = tempfile.mkstemp(suffix=".bz2", dir=scratch_dir)
    try:
        with os.fdopen(tmp_fd, "wb") as cf:          # buffer the COMPRESSED atom (small vs the tensor)
            try:
                _stream_atom_raw(bz2_key, cf)
            except Exception as be:
                if _is_not_found(be):
                    raise raw_miss                   # neither raw nor .bz2 → the ORIGINAL miss propagates
                raise
        try:                                         # clear any partial bytes before decompressing in
            fileobj.seek(0); fileobj.truncate(0)
        except Exception:
            pass
        _decompress_bz2_stream(tmp_path, fileobj)    # STREAMING decompress → never buffers the tensor
    finally:
        try:
            os.unlink(tmp_path)
        except Exception:
            pass


def write_meta(context_uuid: str, doc_id: str, meta: dict) -> None:
    """Write permanent document metadata. Always uncompressed."""
    key = meta_key(context_uuid, doc_id)
    data = json.dumps(meta, ensure_ascii=False, indent=2).encode("utf-8")
    get_primary_store().put_bytes(key, data, content_type="application/json")


# ── Message helpers ────────────────────────────────────────────────────────────

def make_message(doc_id: str, stage: str, storage_key_path: str,
                 context_uuid: str, metadata: dict) -> dict:
    return {
        "doc_id":       doc_id,
        "stage":        stage,
        "storage_key":  storage_key_path,
        "context_uuid": context_uuid,
        "metadata":     json.dumps(metadata),
    }


def parse_message(fields: dict) -> tuple[str, str, str, str, dict]:
    """Returns (doc_id, stage, storage_key, context_uuid, metadata)."""
    return (
        fields["doc_id"],
        fields["stage"],
        fields["storage_key"],
        fields["context_uuid"],
        json.loads(fields.get("metadata", "{}")),
    )


def publish(stage: str, doc_id: str, storage_key_path: str,
            context_uuid: str, metadata: dict,
            r: Optional[_redis_module.Redis] = None) -> None:
    """Publish a stage completion message to the next stream."""
    r = r or get_redis()
    stream = STREAMS[stage]
    msg = make_message(doc_id, stage, storage_key_path, context_uuid, metadata)
    r.xadd(stream, msg)
    logger.info("Published → %s  doc=%s", stream, doc_id[:8])


# ── Base worker ────────────────────────────────────────────────────────────────

class StageWorker:
    """
    Base class for all pipeline stage workers.

    Subclasses implement process() which receives the parsed message and
    the loaded payload dict, and returns the output payload dict.

    The base class handles:
    - Redis stream consumption with consumer group
    - Checkpoint detection (skip if output already exists)
    - Garage read/write
    - Publishing to next stream
    - Error handling and acknowledgment
    """
    input_stage:  str   # stream to read from
    output_stage: str   # stream to publish to after processing
    consumer_name: str  # unique per container instance

    def __init__(self):
        self.r = get_redis()
        ensure_streams(self.r)
        self._running = True
        import signal
        signal.signal(signal.SIGTERM, self._stop)
        signal.signal(signal.SIGINT,  self._stop)

    def _stop(self, sig, _):
        logger.info("Signal %d — stopping after current message.", sig)
        self._running = False

    def process(self, doc_id: str, context_uuid: str,
                metadata: dict, payload: Any) -> Any:
        """Override in subclass. Return output payload."""
        raise NotImplementedError

    def run(self):
        stream    = STREAMS[self.input_stage]
        group     = CONSUMER_GROUPS[stream]
        logger.info("%s starting — stream=%s consumer=%s",
                    self.__class__.__name__, stream, self.consumer_name)

        while self._running:
            try:
                messages = self.r.xreadgroup(
                    group, self.consumer_name,
                    {stream: ">"},
                    count=1,
                    block=5000,
                )
            except Exception as exc:
                logger.error("Redis read error: %s — retrying in 5s", exc)
                time.sleep(5)
                continue

            if not messages:
                continue

            _, entries = messages[0]
            if not entries:
                continue

            msg_id, fields = entries[0]
            doc_id, stage, in_key, context_uuid, metadata = parse_message(fields)

            try:
                self._handle(msg_id, group, stream, doc_id, stage,
                             in_key, context_uuid, metadata)
            except Exception as exc:
                logger.error("Stage %s failed for doc %s: %s",
                             self.input_stage, doc_id[:8], exc, exc_info=True)
                # Ack anyway — don't block the queue on a persistent error
                # The stage can be retried by resubmitting the doc
                self.r.xack(stream, group, msg_id)

        logger.info("%s stopped.", self.__class__.__name__)

    def _handle(self, msg_id, group, stream, doc_id, stage,
                in_key, context_uuid, metadata):
        # Checkpoint: if output already exists, skip processing
        if stage_exists(context_uuid, doc_id, self.output_stage):
            logger.info("Checkpoint hit: %s/%s already done — skipping forward",
                        doc_id[:8], self.output_stage)
            out_key = storage_key(context_uuid, doc_id, self.output_stage)
            publish(self.output_stage, doc_id, out_key, context_uuid, metadata)
            self.r.xack(stream, group, msg_id)
            return

        # Load input payload from Garage
        payload = read_stage(in_key)

        t0 = time.perf_counter()
        output = self.process(doc_id, context_uuid, metadata, payload)
        elapsed = time.perf_counter() - t0

        # Write output to Garage
        out_key = write_stage(context_uuid, doc_id, self.output_stage, output)
        logger.info("%s: doc=%s  %.1fs → %s",
                    self.__class__.__name__, doc_id[:8], elapsed, self.output_stage)

        # Publish next stage
        publish(self.output_stage, doc_id, out_key, context_uuid, metadata)

        # Ack
        self.r.xack(stream, group, msg_id)
