src/proofframe/storage.py

"""Storage backends."""

from __future__ import annotations

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

from .checksum import sha256_hex
from .config import ConfigurationError, Settings
from .models import CampaignManifest, GeneratedMedia, StoredObject


class StorageBackend(Protocol):
    name: str

    def put_media(self, campaign_id: str, media: GeneratedMedia) -> StoredObject: ...

    def put_manifest(self, campaign_id: str, manifest: CampaignManifest) -> StoredObject: ...


def _manifest_bytes(manifest: CampaignManifest) -> bytes:
    return manifest.model_dump_json(indent=2).encode("utf-8")


def _endpoint_with_scheme(endpoint_url: str) -> str:
    endpoint = endpoint_url.strip()
    if endpoint and "://" not in endpoint:
        return f"https://{endpoint}"
    return endpoint


class LocalStorageBackend:
    """Filesystem storage used for tests and credential-free demos."""

    name = "local"

    def __init__(self, root: Path | str = "var/storage") -> None:
        self.root = Path(root)

    def put_media(self, campaign_id: str, media: GeneratedMedia) -> StoredObject:
        return self.put_bytes(
            campaign_id=campaign_id,
            filename=media.filename,
            data=media.data,
            content_type=media.content_type,
        )

    def put_manifest(self, campaign_id: str, manifest: CampaignManifest) -> StoredObject:
        return self.put_bytes(
            campaign_id=campaign_id,
            filename=f"{campaign_id}-manifest.json",
            data=_manifest_bytes(manifest),
            content_type="application/json",
        )

    def put_bytes(
        self, campaign_id: str, filename: str, data: bytes, content_type: str
    ) -> StoredObject:
        del content_type
        folder = self.root / campaign_id
        folder.mkdir(parents=True, exist_ok=True)
        path = folder / filename
        path.write_bytes(data)
        relative = path.relative_to(self.root)
        return StoredObject(
            storage_backend=self.name,
            storage_key=str(relative),
            public_url=f"/storage/{relative.as_posix()}",
            sha256=sha256_hex(data),
            bytes_size=len(data),
        )

    def read_bytes(self, storage_key: str) -> bytes:
        path = (self.root / storage_key).resolve()
        root = self.root.resolve()
        if root not in path.parents:
            raise ConfigurationError(f"Refusing to read outside storage root: {storage_key}")
        return path.read_bytes()


@dataclass
class B2StorageBackend:
    """Backblaze B2 S3-compatible storage backend.

    The backend is intentionally explicit: if `PROOFFRAME_STORAGE_BACKEND=b2` is used,
    missing config raises before any silent local fallback can happen.
    """

    endpoint_url: str
    bucket: str
    key_id: str
    application_key: str
    public_base_url: str = ""
    client: Any | None = None

    name = "b2"

    def __post_init__(self) -> None:
        self.endpoint_url = _endpoint_with_scheme(self.endpoint_url)

    @classmethod
    def from_settings(cls, settings: Settings) -> "B2StorageBackend":
        settings.require_b2()
        return cls(
            endpoint_url=settings.b2_endpoint_url,
            bucket=settings.b2_bucket,
            key_id=settings.b2_key_id,
            application_key=settings.b2_application_key,
            public_base_url=settings.b2_public_base_url,
        )

    def _client(self) -> Any:
        if self.client is not None:
            return self.client
        try:
            import boto3  # type: ignore[import-not-found]
        except ModuleNotFoundError as exc:
            raise ConfigurationError(
                "B2 storage requires boto3. Install with `pip install -e '.[integrations]'`."
            ) from exc
        self.client = boto3.client(
            "s3",
            endpoint_url=self.endpoint_url,
            aws_access_key_id=self.key_id,
            aws_secret_access_key=self.application_key,
        )
        return self.client

    def put_media(self, campaign_id: str, media: GeneratedMedia) -> StoredObject:
        return self.put_bytes(
            key=f"campaigns/{campaign_id}/media/{media.filename}",
            data=media.data,
            content_type=media.content_type,
            metadata={
                "campaign-id": campaign_id,
                "provider": media.provider,
                "model": media.model,
                "sha256": sha256_hex(media.data),
            },
        )

    def put_manifest(self, campaign_id: str, manifest: CampaignManifest) -> StoredObject:
        data = _manifest_bytes(manifest)
        return self.put_bytes(
            key=f"campaigns/{campaign_id}/manifests/{campaign_id}-manifest.json",
            data=data,
            content_type="application/json",
            metadata={"campaign-id": campaign_id, "sha256": sha256_hex(data)},
        )

    def put_bytes(
        self, key: str, data: bytes, content_type: str, metadata: dict[str, str] | None = None
    ) -> StoredObject:
        checksum = sha256_hex(data)
        self._client().put_object(
            Bucket=self.bucket,
            Key=key,
            Body=data,
            ContentType=content_type,
            Metadata=metadata or {"sha256": checksum},
        )
        public_url = f"{self.public_base_url.rstrip('/')}/{key}" if self.public_base_url else None
        return StoredObject(
            storage_backend=self.name,
            storage_key=key,
            public_url=public_url,
            sha256=checksum,
            bytes_size=len(data),
        )


def create_storage_backend(settings: Settings) -> StorageBackend:
    if settings.storage_backend == "local":
        return LocalStorageBackend(settings.storage_root)
    if settings.storage_backend == "b2":
        return B2StorageBackend.from_settings(settings)
    raise ConfigurationError(f"Unknown storage backend: {settings.storage_backend}")


def manifest_to_plain_json(manifest: CampaignManifest) -> dict[str, Any]:
    return json.loads(manifest.model_dump_json())