src/storage.py

"""Storage abstraction layer for garment images and data.

Supports two backends controlled by STORAGE_BACKEND env var:
- "local" (default): filesystem at data/garments/
- "s3": HuggingFace Spaces S3 bucket for persistent storage
"""

import io
import logging
import os
from pathlib import Path

logger = logging.getLogger(__name__)

STORAGE_BACKEND = os.environ.get("STORAGE_BACKEND", "local")
LOCAL_GARMENTS_DIR = Path(__file__).resolve().parent.parent / "data" / "garments"

S3_BUCKET_NAME = os.environ.get("S3_BUCKET_NAME", "")
S3_ENDPOINT_URL = os.environ.get("S3_ENDPOINT_URL", "")
S3_PREFIX = os.environ.get("S3_PREFIX", "garments/")

_s3_client = None


def _get_s3_client():
    global _s3_client
    if _s3_client is None:
        import boto3

        kwargs = {}
        if S3_ENDPOINT_URL:
            kwargs["endpoint_url"] = S3_ENDPOINT_URL

        _s3_client = boto3.client("s3", **kwargs)
    return _s3_client


def _ensure_local_dir():
    LOCAL_GARMENTS_DIR.mkdir(parents=True, exist_ok=True)


def save_image(garment_id: str, image_bytes: bytes, extension: str = "jpg") -> str:
    """Save a garment image and return its storage reference.

    The reference is backend-agnostic (just the filename). Use
    get_image_url() to resolve it to a serveable path/URL.
    """
    filename = f"{garment_id}.{extension}"

    if STORAGE_BACKEND == "s3":
        _save_to_s3(filename, image_bytes)
    else:
        _save_to_local(filename, image_bytes)

    logger.info("Saved image: %s (backend: %s)", filename, STORAGE_BACKEND)
    return filename


def _save_to_local(filename: str, image_bytes: bytes):
    _ensure_local_dir()
    path = LOCAL_GARMENTS_DIR / filename
    path.write_bytes(image_bytes)


def _save_to_s3(filename: str, image_bytes: bytes):
    client = _get_s3_client()
    key = f"{S3_PREFIX}{filename}"
    client.put_object(
        Bucket=S3_BUCKET_NAME,
        Key=key,
        Body=image_bytes,
        ContentType="image/jpeg",
    )


def get_image_path(garment_id: str, extension: str = "jpg") -> str | None:
    """Return the local filesystem path for a garment image.

    For S3 backend, downloads to a local cache first.
    Returns None if the image doesn't exist.
    """
    filename = f"{garment_id}.{extension}"

    if STORAGE_BACKEND == "s3":
        return _download_from_s3(filename)

    path = LOCAL_GARMENTS_DIR / filename
    if path.exists():
        return str(path)
    return None


def _download_from_s3(filename: str) -> str | None:
    _ensure_local_dir()
    local_path = LOCAL_GARMENTS_DIR / filename

    if local_path.exists():
        return str(local_path)

    try:
        client = _get_s3_client()
        key = f"{S3_PREFIX}{filename}"
        response = client.get_object(Bucket=S3_BUCKET_NAME, Key=key)
        local_path.write_bytes(response["Body"].read())
        return str(local_path)
    except Exception as e:
        logger.warning("Failed to download %s from S3: %s", filename, e)
        return None


def delete_image(garment_id: str, extension: str = "jpg"):
    """Remove a garment image from storage."""
    filename = f"{garment_id}.{extension}"

    if STORAGE_BACKEND == "s3":
        try:
            client = _get_s3_client()
            key = f"{S3_PREFIX}{filename}"
            client.delete_object(Bucket=S3_BUCKET_NAME, Key=key)
        except Exception as e:
            logger.warning("Failed to delete %s from S3: %s", filename, e)

    local_path = LOCAL_GARMENTS_DIR / filename
    if local_path.exists():
        local_path.unlink()


def image_exists(garment_id: str, extension: str = "jpg") -> bool:
    """Check if a garment image exists in storage."""
    filename = f"{garment_id}.{extension}"

    if STORAGE_BACKEND == "s3":
        try:
            client = _get_s3_client()
            key = f"{S3_PREFIX}{filename}"
            client.head_object(Bucket=S3_BUCKET_NAME, Key=key)
            return True
        except Exception:
            return False

    return (LOCAL_GARMENTS_DIR / filename).exists()