"""Storage backend implementations.

All blobs (objects) are content-addressed by ``object_id``.  The only backend
is ``BlobBackend``, which works against any S3-compatible store (Cloudflare R2,
MinIO, AWS S3).  Configure via ``BLOB_STORAGE_BUCKET`` / ``BLOB_STORAGE_ENDPOINT``
env vars.

URI scheme:
    ``s3://<bucket>/objects/<object_id>``  — AWS S3 / S3-compatible
"""

import base64
import concurrent.futures
import logging
from collections.abc import AsyncIterator
from typing import TYPE_CHECKING, Protocol

if TYPE_CHECKING:
    from musehub.db.musehub_repo_models import MusehubObject

from musehub.config import settings
from musehub.types.json_types import StrDict

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# S3 concurrency constants
# ---------------------------------------------------------------------------

# Size of the dedicated thread pool used for S3/R2 blocking I/O.
# boto3 is not async-native, so every put/get/delete goes through
# asyncio.get_running_loop().run_in_executor().  The default executor has
# min(32, cpu+4) threads — on a 4-CPU server that's 8, which throttles
# concurrent R2 puts to 8 even when the asyncio semaphore allows 100.
# Using a dedicated pool with S3_THREAD_POOL_SIZE workers lets us saturate
# the R2 connection pool (S3_POOL_CONNECTIONS) without queuing in the
# scheduler.
S3_THREAD_POOL_SIZE: int = 100

# Maximum simultaneous TCP connections per boto3 S3 client.
# The default (10) serializes 100 concurrent asyncio-to-thread dispatches
# through a 10-slot queue.  100 connections matches the thread pool size so
# every in-flight thread can make progress immediately.
S3_POOL_CONNECTIONS: int = 100

# Lazy-initialized thread pool — created on first S3 put, shared for the
# lifetime of the process.
_s3_thread_pool: concurrent.futures.ThreadPoolExecutor | None = None

def _get_s3_thread_pool() -> concurrent.futures.ThreadPoolExecutor:
    global _s3_thread_pool
    if _s3_thread_pool is None:
        _s3_thread_pool = concurrent.futures.ThreadPoolExecutor(
            max_workers=S3_THREAD_POOL_SIZE,
            thread_name_prefix="muse-s3",
        )
    return _s3_thread_pool

# Named type for the object map returned by get_batch.
type ObjectMap = dict[str, bytes]

# ── Boto3 client Protocol ─────────────────────────────────────────────────────

class _StreamingBody(Protocol):
    def read(self, size: int = -1) -> bytes: ...

class _S3Response(Protocol):
    def __getitem__(self, key: str) -> _StreamingBody: ...

class _S3Client(Protocol):
    def head_object(self, *, Bucket: str, Key: str) -> _S3Response: ...
    def put_object(self, *, Bucket: str, Key: str, Body: bytes) -> _S3Response: ...
    def get_object(self, *, Bucket: str, Key: str) -> _S3Response: ...
    def delete_object(self, *, Bucket: str, Key: str) -> _S3Response: ...
    def generate_presigned_url(
        self,
        ClientMethod: str,
        Params: StrDict,
        ExpiresIn: int,
    ) -> str: ...

# ── Storage backend Protocol ──────────────────────────────────────────────────

class StorageBackend(Protocol):
    """Content-addressed binary store protocol.

    Storage is globally content-addressed — every method operates on
    ``object_id`` alone; there is no per-repo namespace.

    Backends must be safe to call from async code.  I/O-bound implementations
    should use ``asyncio.to_thread`` internally rather than blocking the event
    loop.
    """

    async def put(self, object_id: str, data: bytes) -> str:
        """Persist ``data`` and return a ``storage_uri``."""
        ...

    async def get(self, object_id: str) -> bytes | None:
        """Return raw bytes for the object, or ``None`` if not found."""
        ...

    async def get_batch(self, object_ids: list[str]) -> ObjectMap:
        """Return ``{object_id: bytes}`` for all *object_ids* that exist.

        Default: sequential fallback via ``get()``.  Implementations should
        override this to do all reads in a single thread dispatch.
        """
        result: ObjectMap = {}
        for oid in object_ids:
            data = await self.get(oid)
            if data is not None:
                result[oid] = data
        return result

    async def exists(self, object_id: str) -> bool:
        """Return True if the object is already stored (avoids re-upload)."""
        ...

    async def delete(self, object_id: str) -> None:
        """Remove an object (used on repo deletion)."""
        ...

    def uri_for(self, object_id: str) -> str:
        """Return the canonical storage URI without necessarily persisting."""
        ...

    async def stream(
        self, object_id: str, chunk_size: int = 65536
    ) -> AsyncIterator[bytes]:
        """Yield the object's bytes in chunks without buffering the full file.

        Implementations must yield an empty iterator (not raise) when the object
        does not exist.  Callers should check ``exists()`` first if a 404 is
        needed; this method is optimised for the streaming-download hot path.
        """
        ...  # pragma: no cover
        yield b""  # satisfy type checker — real impls override this

    async def presign_batch(
        self,
        object_ids: list[str],
        direction: str,
        ttl_seconds: int,
    ) -> StrDict:
        """Return presigned URLs keyed by object_id.  Default: empty (not supported)."""
        return {}

class BlobBackend:
    """AWS S3 (or S3-compatible) storage backend.

    Requires ``boto3`` to be installed and AWS credentials to be configured
    (via environment variables, IAM instance profile, or ``~/.aws/credentials``).

    Objects are stored at:
        ``s3://<bucket>/objects/<object_id>``

    Storage is globally content-addressed — identical bytes are stored once
    regardless of which repo pushed them.

    Uploads are idempotent — ``put`` calls PutObject directly without a prior
    HeadObject check.  Deduplication is handled at the database layer via the
    ``filter-objects`` endpoint: only objects missing from the DB reach this
    method.  Skipping HeadObject halves the R2 round-trips per object and
    doubles effective throughput.  R2 PutObject is idempotent (same content →
    same key → no-op at the storage layer), so the double-write case is safe.
    """

    supports_presign: bool = True

    def __init__(
        self,
        bucket: str | None = None,
        region: str | None = None,
        endpoint_url: str | None = None,
        access_key_id: str | None = None,
        secret_access_key: str | None = None,
        public_endpoint_url: str | None = None,
    ) -> None:
        self._bucket = bucket or settings.blob_storage_bucket or settings.aws_s3_asset_bucket or ""
        self._region = region or (settings.blob_storage_region if settings.blob_storage_bucket else settings.aws_region)
        self._endpoint_url = endpoint_url or settings.blob_storage_endpoint
        self._public_endpoint_url = public_endpoint_url or settings.blob_storage_public_endpoint
        self._cdn_base_url = settings.blob_storage_cdn_base_url
        self._access_key_id = access_key_id or settings.blob_storage_access_key_id
        self._secret_access_key = secret_access_key or settings.blob_storage_secret_access_key
        self._client: _S3Client | None = None

    def _get_client(self) -> _S3Client:
        if self._client is None:
            import boto3
            from botocore.config import Config
            kwargs = {
                "region_name": self._region or "",
                # Increase the connection pool beyond the 10-connection default.
                # With S3_THREAD_POOL_SIZE worker threads all making concurrent
                # put_object calls, a 10-connection pool serializes 90 % of them
                # into a wait queue.  S3_POOL_CONNECTIONS connections lets every
                # in-flight thread make immediate progress.
                "config": Config(
                    max_pool_connections=S3_POOL_CONNECTIONS,
                    # Suppress optional CRC32 checksums — MinIO echoes them back,
                    # adding response header bytes. Keeping these off reduces wire
                    # overhead and keeps the 200 OK response compact.
                    request_checksum_calculation="when_required",
                    response_checksum_validation="when_required",
                ),
            }
            if self._endpoint_url:
                kwargs["endpoint_url"] = self._endpoint_url
            if self._access_key_id and self._secret_access_key:
                kwargs["aws_access_key_id"] = self._access_key_id
                kwargs["aws_secret_access_key"] = self._secret_access_key
            self._client = boto3.client("s3", **kwargs)
            # Python 3.14 regression: when botocore sends Expect: 100-continue,
            # MinIO sends 100 Continue then 200 OK in the same TCP segment.
            # Python 3.14's http.client sees the 200 OK bytes as unparsed data
            # after the 100 Continue headers → MissingHeaderBodySeparatorDefect.
            # urllib3 catches the defect but returns the connection to the pool
            # without draining the response body (dirty connection). The next PUT
            # on the same pool reuses this connection; MinIO reads the leftover
            # bytes as a new request, becomes confused, and closes the connection
            # after its 30s keep-alive timeout. Fix: remove Expect: 100-continue
            # so botocore sends headers+body in one shot → single 200 OK → clean.
            def _remove_expect_continue(request: "botocore.awsrequest.AWSPreparedRequest", **kwargs: str) -> None:
                request.headers.pop("Expect", None)
            self._client.meta.events.register(
                "before-send.s3.PutObject",
                _remove_expect_continue,
            )
        return self._client

    def _key(self, object_id: str) -> str:
        # object_id is canonical algo:hex (e.g. "sha256:<64-hex>").
        # Colons are valid in S3/R2 keys — no substitution needed.
        return f"objects/{object_id}"

    def uri_for(self, object_id: str) -> str:
        return f"s3://{self._bucket}/{self._key(object_id)}"

    async def put(self, object_id: str, data: bytes) -> str:
        import asyncio
        key = self._key(object_id)
        import logging as _log
        _log.getLogger(__name__).debug("BlobBackend.put: oid=%s key=%s size=%d", object_id, key, len(data))
        client = self._get_client()
        loop = asyncio.get_running_loop()
        await loop.run_in_executor(_get_s3_thread_pool(), self._s3_put, client, key, data)
        return self.uri_for(object_id)

    def _s3_put(self, client: _S3Client, key: str, data: bytes) -> None:
        # No HeadObject check — deduplication is handled upstream by the
        # filter-objects endpoint.  Only objects confirmed missing from the DB
        # reach this call, so the check is redundant and wastes one R2
        # round-trip per object.  PutObject is idempotent: uploading the same
        # content-addressed bytes twice is a safe no-op at the storage layer.
        client.put_object(Bucket=self._bucket, Key=key, Body=data)

    async def get(self, object_id: str) -> bytes | None:
        import asyncio
        c = self._get_client()
        key = self._key(object_id)
        try:
            def _do_get() -> _S3Response:
                return c.get_object(Bucket=self._bucket, Key=key)
            result = await asyncio.to_thread(_do_get)
            data: bytes = result["Body"].read()
            return data
        except Exception:
            return None

    async def get_batch(self, object_ids: list[str]) -> ObjectMap:
        """Fetch all objects in parallel via asyncio.gather.

        The base-class fallback is sequential — one S3 round-trip per object.
        On R2/S3 in production that means a 100+ object clone hits Cloudflare's
        100-second origin timeout. Overriding here dispatches all fetches
        concurrently so total latency ≈ max(single object latency).
        """
        import asyncio

        results = await asyncio.gather(*(self.get(oid) for oid in object_ids))
        return {
            oid: data
            for oid, data in zip(object_ids, results)
            if data is not None
        }

    async def exists(self, object_id: str) -> bool:
        import asyncio
        c = self._get_client()
        key = self._key(object_id)

        def _head() -> bool:
            try:
                c.head_object(Bucket=self._bucket, Key=key)
                return True
            except Exception:
                return False

        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(_get_s3_thread_pool(), _head)

    async def delete(self, object_id: str) -> None:
        import asyncio
        key = self._key(object_id)
        c = self._get_client()
        def _do_delete() -> None:
            c.delete_object(Bucket=self._bucket, Key=key)
        await asyncio.to_thread(_do_delete)

    async def stream(
        self, object_id: str, chunk_size: int = 65536
    ) -> AsyncIterator[bytes]:
        """Yield S3 object content in ``chunk_size`` byte chunks.

        Uses boto3's streaming body so the full object is never loaded into
        memory — the HTTP response body is consumed incrementally.
        """
        import asyncio

        c = self._get_client()
        key = self._key(object_id)
        body: _StreamingBody | None = None
        try:
            def _get_stream() -> _StreamingBody:
                return c.get_object(Bucket=self._bucket, Key=key)["Body"]
            body = await asyncio.to_thread(_get_stream)
        except Exception:
            return  # object not found — yield nothing

        assert body is not None
        def _read_chunk() -> bytes:
            return body.read(chunk_size)

        while True:
            chunk = await asyncio.to_thread(_read_chunk)
            if not chunk:
                break
            yield chunk

    def _rewrite_presign_url(self, url: str) -> str:
        """Replace the internal endpoint host with the public endpoint host.

        boto3 embeds whatever endpoint_url the client was initialised with into
        every presigned URL.  In local dev that's http://minio:9000 (Docker
        internal), which is unreachable from the host machine.  When
        BLOB_STORAGE_PUBLIC_ENDPOINT is set, swap in that host so external clients can
        PUT/GET directly.
        """
        if not self._public_endpoint_url or not self._endpoint_url:
            return url
        return url.replace(self._endpoint_url, self._public_endpoint_url, 1)

    def _mpack_key(self, mpack_key: str) -> str:
        return f"mpacks/{mpack_key}"

    async def presign_mpack_put(self, mpack_key: str, ttl_seconds: int = 3600) -> str:
        """Return a presigned PUT URL for an mpack blob (stored under mpacks/ prefix)."""
        import asyncio
        c = self._get_client()
        key = self._mpack_key(mpack_key)
        logger.warning("[presign_mpack_put] mpack_key=%s  s3_key=%s  bucket=%s", mpack_key, key, self._bucket)

        def _presign() -> str:
            raw_url = c.generate_presigned_url(
                "put_object",
                Params={"Bucket": self._bucket, "Key": key, "ContentType": "application/x-muse-pack"},
                ExpiresIn=ttl_seconds,
            )
            logger.warning("[presign_mpack_put] raw_url (before rewrite) = %s", raw_url)
            return raw_url

        loop = asyncio.get_running_loop()
        final_url = self._rewrite_presign_url(await loop.run_in_executor(_get_s3_thread_pool(), _presign))
        logger.warning("[presign_mpack_put] final_url (after rewrite)  = %s", final_url)
        return final_url

    async def presign_mpack_get(self, mpack_key: str, ttl_seconds: int = 3600) -> str:
        """Return a presigned GET URL for an mpack blob (stored under mpacks/ prefix).

        When ``blob_storage_cdn_base_url`` is configured the presigned URL host is
        replaced with the CDN base URL so repeat clones hit the Cloudflare edge
        cache instead of R2/MinIO directly.
        """
        import asyncio
        c = self._get_client()
        key = self._mpack_key(mpack_key)

        def _presign() -> str:
            return c.generate_presigned_url(
                "get_object",
                Params={"Bucket": self._bucket, "Key": key},
                ExpiresIn=ttl_seconds,
            )

        loop = asyncio.get_running_loop()
        url = self._rewrite_presign_url(await loop.run_in_executor(_get_s3_thread_pool(), _presign))
        if self._cdn_base_url:
            from urllib.parse import urlparse, urlunparse
            parsed = urlparse(url)
            cdn = urlparse(self._cdn_base_url)
            url = urlunparse(parsed._replace(scheme=cdn.scheme, netloc=cdn.netloc))
        return url

    async def get_range(self, mpack_key: str, byte_offset: int, byte_length: int) -> bytes | None:
        """Fetch a byte range from an mpack using an S3 Range GET.

        Issues ``Range: bytes=<byte_offset>-<byte_offset+byte_length-1>`` so
        only the requested slice is transferred — O(object_size) instead of
        O(mpack_size).  Returns ``None`` if the mpack or range is not found.
        """
        import asyncio
        c = self._get_client()
        key = self._mpack_key(mpack_key)
        range_header = f"bytes={byte_offset}-{byte_offset + byte_length - 1}"

        def _get() -> bytes | None:
            try:
                resp = c.get_object(Bucket=self._bucket, Key=key, Range=range_header)
                return resp["Body"].read()
            except Exception:
                return None

        return await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _get)

    async def exists_mpack(self, mpack_key: str) -> bool:
        """Return True if the mpack exists in storage."""
        import asyncio
        c = self._get_client()
        key = self._mpack_key(mpack_key)

        def _head() -> bool:
            try:
                c.head_object(Bucket=self._bucket, Key=key)
                return True
            except Exception:
                return False

        return await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _head)

    async def get_mpack(self, mpack_key: str) -> bytes | None:
        """Fetch the raw mpack bytes from the mpacks/ prefix."""
        import asyncio
        import hashlib as _hashlib
        c = self._get_client()
        key = self._mpack_key(mpack_key)
        logger.warning("[get_mpack] mpack_key=%s  s3_key=%s  bucket=%s", mpack_key, key, self._bucket)

        def _get() -> bytes | None:
            try:
                resp = c.get_object(Bucket=self._bucket, Key=key)
                etag = resp.get("ETag", "")
                content_length = resp.get("ContentLength", "?")
                data = resp["Body"].read()
                actual_sha256 = "sha256:" + _hashlib.sha256(data).hexdigest()
                logger.warning(
                    "[get_mpack] fetched %d bytes  etag=%s  content_length=%s  sha256(data)=%s  matches_key=%s",
                    len(data), etag, content_length, actual_sha256, actual_sha256 == mpack_key,
                )
                return data
            except Exception as exc:
                logger.warning("[get_mpack] FAILED key=%s exc=%s", key, exc)
                return None

        return await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _get)

    async def put_mpack(self, mpack_key: str, data: bytes) -> None:
        """Write raw mpack bytes directly under the mpacks/ prefix.

        Mpacks are content-addressed (key == sha256(bytes)) so the key never
        changes after creation.  ``Cache-Control: public, max-age=31536000, immutable``
        tells Cloudflare it is safe to cache the response indefinitely at the edge.
        """
        import asyncio
        c = self._get_client()
        key = self._mpack_key(mpack_key)

        def _put() -> None:
            c.put_object(
                Bucket=self._bucket,
                Key=key,
                Body=data,
                ContentType="application/x-muse-pack",
                CacheControl="public, max-age=31536000, immutable",
            )

        await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _put)

    async def quarantine_mpack(self, mpack_key: str) -> None:
        """Move an mpack blob from mpacks/ to quarantine/ prefix.

        Called by process_mpack_index_job on MPackValidationError to
        isolate suspect content from normal mpack storage.
        """
        import asyncio
        c = self._get_client()
        src_key = self._mpack_key(mpack_key)
        dst_key = f"quarantine/{mpack_key}"

        def _move() -> None:
            try:
                c.copy_object(
                    Bucket=self._bucket,
                    CopySource={"Bucket": self._bucket, "Key": src_key},
                    Key=dst_key,
                )
                c.delete_object(Bucket=self._bucket, Key=src_key)
            except Exception:
                pass  # best-effort; do not mask the validation error

        await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _move)

    async def quarantine_object(self, object_id: str) -> None:
        """Move an object blob from objects/ to quarantine/ prefix.

        Called by the DMCA takedown endpoint to remove a blob from public access.
        Best-effort: silently no-ops if the object is not present.
        """
        import asyncio
        c = self._get_client()
        src_key = self._key(object_id)
        dst_key = f"quarantine/{object_id}"

        def _move() -> None:
            try:
                c.copy_object(
                    Bucket=self._bucket,
                    CopySource={"Bucket": self._bucket, "Key": src_key},
                    Key=dst_key,
                )
                c.delete_object(Bucket=self._bucket, Key=src_key)
            except Exception:
                pass

        await asyncio.get_running_loop().run_in_executor(_get_s3_thread_pool(), _move)

    async def presign_put(self, object_id: str, ttl_seconds: int = 3600) -> str:
        """Return a presigned PUT URL for *object_id* valid for *ttl_seconds*.

        The client PUTs raw bytes directly to R2 using this URL, bypassing the
        API server and Cloudflare entirely.  After all uploads complete, the
        client calls ``POST /push/objects/confirm`` to register the objects in
        the MuseHub DB.
        """
        import asyncio

        c = self._get_client()
        key = self._key(object_id)

        def _presign() -> str:
            return c.generate_presigned_url(
                "put_object",
                Params={"Bucket": self._bucket, "Key": key, "ContentType": "application/octet-stream"},
                ExpiresIn=ttl_seconds,
            )

        loop = asyncio.get_running_loop()
        return self._rewrite_presign_url(await loop.run_in_executor(_get_s3_thread_pool(), _presign))

    async def presign_get(self, object_id: str, ttl_seconds: int = 3600) -> str:
        """Return a presigned GET URL for *object_id* valid for *ttl_seconds*.

        The client GETs raw bytes directly from R2 using this URL, bypassing
        the API server.  Used in the fetch path when the backend is R2.
        """
        import asyncio

        c = self._get_client()
        key = self._key(object_id)

        def _presign() -> str:
            return c.generate_presigned_url(
                "get_object",
                Params={"Bucket": self._bucket, "Key": key},
                ExpiresIn=ttl_seconds,
            )

        loop = asyncio.get_running_loop()
        return self._rewrite_presign_url(
            await loop.run_in_executor(_get_s3_thread_pool(), _presign)
        )

    async def presign_batch(
        self,
        object_ids: list[str],
        direction: str,
        ttl_seconds: int,
    ) -> StrDict:
        """Return presigned URLs for *object_ids* in a single thread dispatch."""
        import asyncio

        c = self._get_client()
        method = "put_object" if direction == "put" else "get_object"

        def _presign_all() -> list[str]:
            return [
                c.generate_presigned_url(
                    method,
                    Params={"Bucket": self._bucket, "Key": self._key(oid)},
                    ExpiresIn=ttl_seconds,
                )
                for oid in object_ids
            ]

        urls = await asyncio.to_thread(_presign_all)
        return {oid: self._rewrite_presign_url(url) for oid, url in zip(object_ids, urls)}

class MemoryBackend:
    """In-memory blob backend for tests — no external services required."""

    supports_presign: bool = False

    def __init__(self) -> None:
        self._store: dict[str, bytes] = {}
        self._mpacks: dict[str, bytes] = {}

    def uri_for(self, object_id: str) -> str:
        # Use s3:// prefix so read_object_bytes routes through get_backend().get(),
        # which MemoryBackend implements in-memory without external services.
        return f"s3://memory-bucket/objects/{object_id}"

    async def put(self, object_id: str, data: bytes) -> str:
        self._store[object_id] = data
        return self.uri_for(object_id)

    async def get(self, object_id: str) -> bytes | None:
        return self._store.get(object_id)

    async def get_batch(self, object_ids: list[str]) -> "ObjectMap":
        return {oid: self._store[oid] for oid in object_ids if oid in self._store}

    async def exists(self, object_id: str) -> bool:
        return object_id in self._store

    async def delete(self, object_id: str) -> None:
        self._store.pop(object_id, None)

    async def stream(self, object_id: str, chunk_size: int = 65536) -> "AsyncIterator[bytes]":
        data = self._store.get(object_id)
        if data:
            yield data

    async def presign_batch(self, object_ids: list[str], direction: str, ttl_seconds: int) -> "StrDict":
        return {}

    async def presign_get(self, object_id: str, ttl_seconds: int = 3600) -> str:
        return f"memory://objects/{object_id}?op=get&ttl={ttl_seconds}"

    # ── mpack interface ───────────────────────────────────────────────────────

    async def put_mpack(self, mpack_key: str, data: bytes) -> None:
        self._mpacks[mpack_key] = data

    async def get_mpack(self, mpack_key: str) -> bytes | None:
        return self._mpacks.get(mpack_key)

    async def exists_mpack(self, mpack_key: str) -> bool:
        return mpack_key in self._mpacks

    async def get_range(self, mpack_key: str, byte_offset: int, byte_length: int) -> bytes | None:
        data = self._mpacks.get(mpack_key)
        if data is None:
            return None
        return data[byte_offset: byte_offset + byte_length]

    async def presign_mpack_put(self, mpack_key: str, ttl_seconds: int = 3600) -> str:
        return f"memory://mpacks/{mpack_key}?op=put"

    async def presign_mpack_get(self, mpack_key: str, ttl_seconds: int = 3600) -> str:
        return f"memory://mpacks/{mpack_key}?op=get"

    async def quarantine_mpack(self, mpack_key: str) -> None:
        self._mpacks.pop(mpack_key, None)

    async def quarantine_object(self, object_id: str) -> None:
        self._store.pop(object_id, None)


def _get_backend_impl() -> StorageBackend:
    """Real implementation — not patched by conftest. Tests that verify the
    selection logic call this directly."""
    if settings.blob_storage_bucket:
        return BlobBackend()
    if settings.aws_s3_asset_bucket:
        return BlobBackend()
    raise RuntimeError(
        "No storage backend configured. Set BLOB_STORAGE_BUCKET "
        "or AWS_S3_ASSET_BUCKET."
    )


def get_backend() -> StorageBackend:
    """Return the configured BlobBackend."""
    return _get_backend_impl()

async def read_object_bytes(
    obj: "MusehubObject",
    session: "AsyncSession | None" = None,
) -> bytes | None:
    """Return raw bytes for any MusehubObject.

    Resolution order:
        1. ``content_cache`` — in-memory, no I/O
        2. ``storage_uri`` starts with ``s3://`` → BlobBackend.get()
        3. ``storage_uri`` starts with ``mpack://`` → byte-range GET via
           musehub_mpack_index (fast) or full mpack download (slow fallback)
        4. Nothing available → None

    Pass ``session`` when available so byte ranges can be looked up from
    musehub_mpack_index — this avoids downloading the full mpack.
    """
    content_cache = getattr(obj, "content_cache", None)
    if content_cache is not None:
        return content_cache

    storage_uri: str = getattr(obj, "storage_uri", "") or ""
    if storage_uri.startswith("s3://"):
        return await get_backend().get(getattr(obj, "object_id", ""))

    if storage_uri.startswith("mpack://"):
        mpack_id = storage_uri[len("mpack://"):]
        oid: str = getattr(obj, "object_id", "")
        byte_offset: int | None = getattr(obj, "byte_offset", None)
        byte_length: int | None = getattr(obj, "byte_length", None)

        # Look up byte range from musehub_mpack_index when session is available.
        if (byte_offset is None or byte_length is None) and oid and session is not None:
            from musehub.db.musehub_repo_models import MusehubMPackIndex
            from sqlalchemy import select as _sa_select
            # Prefer the specific push mpack this object was stored in.
            _idx = (await session.execute(
                _sa_select(MusehubMPackIndex).where(
                    MusehubMPackIndex.entity_id == oid,
                    MusehubMPackIndex.mpack_id == mpack_id,
                    MusehubMPackIndex.entity_type == "object",
                ).limit(1)
            )).scalar_one_or_none()
            if _idx is not None:
                byte_offset = _idx.byte_offset
                byte_length = _idx.byte_length
            else:
                # The original push mpack may not have index entries (e.g. the
                # prebuilt clone mpack superseded it in the index).  Fall back
                # to any indexed location for this object so we can serve it
                # via range GET without downloading the entire push mpack.
                _any_idx = (await session.execute(
                    _sa_select(MusehubMPackIndex).where(
                        MusehubMPackIndex.entity_id == oid,
                        MusehubMPackIndex.entity_type == "object",
                        MusehubMPackIndex.byte_offset.isnot(None),
                        MusehubMPackIndex.byte_length.isnot(None),
                    ).limit(1)
                )).scalar_one_or_none()
                if _any_idx is not None:
                    mpack_id = _any_idx.mpack_id
                    byte_offset = _any_idx.byte_offset
                    byte_length = _any_idx.byte_length

        # Fast path: byte-range GET — O(object_size), not O(mpack_size).
        if byte_offset is not None and byte_length is not None:
            return await get_backend().get_range(mpack_id, byte_offset, byte_length)

        mpack_raw = await get_backend().get_mpack(mpack_id)
        if mpack_raw is None:
            return None
        import zstandard as _zstd
        _dctx = _zstd.ZstdDecompressor()
        if mpack_raw[:4] != b"MUSE":
            raise ValueError(f"mpack is not MUSE binary format (got {mpack_raw[:4]!r})")
        from muse.core.mpack import parse_wire_mpack as _parse_wire_mpack
        payload = _parse_wire_mpack(mpack_raw)
        for entry in payload.get("blobs", []):
            if entry.get("object_id") == oid:
                content = entry.get("content") or b""
                if not isinstance(content, bytes):
                    content = bytes(content)
                _ZSTD_MAGIC = b"\x28\xb5\x2f\xfd"
                if (entry.get("encoding") == "zstd" or content[:4] == _ZSTD_MAGIC) and content:
                    content = _dctx.decompress(content)
                return content
        return None

    return None

def decode_b64(b64_str: str) -> bytes:
    """Decode a base64 string (with or without padding) into bytes."""
    padded = b64_str + "=" * (-len(b64_str) % 4)
    return base64.b64decode(padded)