storage/artifacts.py

"""
Filesystem artifact store.

Responsible for saving uploaded source images, generated annotated
images, and any other binary artifacts.  Returns relative paths so the
API can construct download URLs without coupling to the filesystem layout.
"""

from __future__ import annotations

import time
from pathlib import Path
from typing import Optional

import cv2
import numpy as np

from loguru import logger


class ArtifactStore:
    """Saves and reads image artifacts under a root directory."""

    def __init__(self, root: Path) -> None:
        self._root = Path(root)
        self._uploads = self._root / "uploads"
        self._generated = self._root / "generated"
        self._uploads.mkdir(parents=True, exist_ok=True)
        self._generated.mkdir(parents=True, exist_ok=True)

    # ------------------------------------------------------------------ #
    # Uploads (source images from users)
    # ------------------------------------------------------------------ #
    def save_upload(self, img: np.ndarray, prefix: str = "upload") -> Path:
        """Save an uploaded image.  Returns absolute path."""
        path = self._uploads / f"{prefix}_{int(time.time() * 1000)}.jpg"
        cv2.imwrite(str(path), img)
        return path

    def save_upload_bytes(self, data: bytes, prefix: str = "upload",
                          suffix: str = ".jpg") -> Path:
        path = self._uploads / f"{prefix}_{int(time.time() * 1000)}{suffix}"
        path.write_bytes(data)
        return path

    # ------------------------------------------------------------------ #
    # Generated artifacts (annotated images, crops, montages)
    # ------------------------------------------------------------------ #
    def save_generated(self, img: np.ndarray, prefix: str = "gen") -> Path:
        path = self._generated / f"{prefix}_{int(time.time() * 1000)}.jpg"
        cv2.imwrite(str(path), img)
        return path

    # ------------------------------------------------------------------ #
    # Lookup
    # ------------------------------------------------------------------ #
    def resolve(self, relative_path: str) -> Optional[Path]:
        """Resolve a relative path to absolute, with root-jail safety."""
        target = (self._root / relative_path).resolve()
        try:
            target.relative_to(self._root.resolve())
        except ValueError:
            logger.warning(f"Path traversal attempt blocked: {relative_path}")
            return None
        if not target.exists():
            return None
        return target

    @property
    def root(self) -> Path:
        return self._root

    @property
    def uploads_dir(self) -> Path:
        return self._uploads

    @property
    def generated_dir(self) -> Path:
        return self._generated