iron/app/services/storage_service.py

from __future__ import annotations

import builtins
import json
from datetime import UTC, datetime
from pathlib import Path

import anyio
from sqlalchemy.ext.asyncio import AsyncSession

from app.adapters.storage.local import LocalStorageAdapter
from app.adapters.storage.s3 import S3StorageAdapter
from app.core.config import get_settings
from app.core.ids import new_id
from app.models.entities import Backend, Blob, BlobReplica, CacheEntry
from app.repositories.backend_repository import BackendRepository
from app.repositories.blob_replica_repository import BlobReplicaRepository
from app.repositories.cache_entry_repository import CacheEntryRepository
from app.services.policy_service import PlacementDecision, PolicyService


class StorageService:
    def __init__(self, session: AsyncSession) -> None:
        self.session = session
        self.settings = get_settings()
        self.backend_repository = BackendRepository(session)
        self.blob_replica_repository = BlobReplicaRepository(session)
        self.cache_entry_repository = CacheEntryRepository(session)
        self.policy_service = PolicyService(session)

    async def ensure_cache_dir(self) -> Path:
        cache_dir = Path(self.settings.cache_dir)
        await anyio.to_thread.run_sync(lambda: cache_dir.mkdir(parents=True, exist_ok=True))
        return cache_dir

    async def get_replication_targets_for_mime_type(
        self,
        mime_type: str | None,
        size_bytes: int | None = None,
    ) -> PlacementDecision:
        return await self.policy_service.get_placement_decision(mime_type, size_bytes=size_bytes)

    async def persist_blob_to_backend(self, blob: Blob, backend: Backend, source_path: str) -> BlobReplica:
        existing = await self.blob_replica_repository.get_by_blob_and_backend(blob.id, backend.id)
        storage_key = self._build_storage_key(blob)
        adapter = self._build_adapter(backend)
        existing_is_usable = False
        if existing is not None and existing.status == "ready":
            try:
                existing_is_usable = await adapter.object_exists(storage_key)
            except Exception:
                existing_is_usable = False
        if existing_is_usable:
            return existing

        if backend.type == "local":
            await adapter.put_file(source_path, storage_key)
            etag = None
        else:
            result = await adapter.put_file(source_path, storage_key)
            etag = result.get("etag")

        replica = existing or BlobReplica(
            id=new_id("rep"),
            blob_id=blob.id,
            backend_id=backend.id,
            storage_key=storage_key,
            size_bytes=blob.size_bytes,
            checksum=blob.content_hash,
            status="ready",
            etag=etag,
            last_verified_at=datetime.now(UTC),
        )
        replica.storage_key = storage_key
        replica.status = "ready"
        replica.size_bytes = blob.size_bytes
        replica.checksum = blob.content_hash
        replica.etag = etag
        replica.last_verified_at = datetime.now(UTC)

        if existing is None:
            created = await self.blob_replica_repository.create(replica)
        else:
            created = await self.blob_replica_repository.save(replica)
        await self.session.commit()
        return created

    async def materialize_replica_to_local_cache(self, blob: Blob, replica: BlobReplica, backend: Backend) -> Path:
        if backend.type == "local":
            adapter = self._build_adapter(backend)
            return await adapter.resolve_path(replica.storage_key)

        cache_dir = await self.ensure_cache_dir()
        cache_path = cache_dir / self._build_storage_key(blob)
        adapter = self._build_adapter(backend)
        await adapter.download_file(replica.storage_key, cache_path)

        existing_entry = await self.cache_entry_repository.get_ready_by_blob(blob.id)
        now = datetime.now(UTC)
        if existing_entry is None:
            entry = CacheEntry(
                id=new_id("cch"),
                blob_id=blob.id,
                local_path=str(cache_path),
                cache_kind="remote-fetch",
                size_bytes=blob.size_bytes,
                status="ready",
                last_accessed_at=now,
                expires_at=None,
            )
            await self.cache_entry_repository.create(entry)
        else:
            existing_entry.local_path = str(cache_path)
            existing_entry.status = "ready"
            existing_entry.size_bytes = blob.size_bytes
            existing_entry.last_accessed_at = now
            await self.cache_entry_repository.save(existing_entry)
        await self.session.commit()
        return cache_path

    async def resolve_best_local_path(self, blob: Blob) -> Path:
        cached_entry = await self.cache_entry_repository.get_ready_by_blob(blob.id)
        if cached_entry is not None:
            cached_path = Path(cached_entry.local_path)
            if await anyio.to_thread.run_sync(cached_path.exists):
                cached_entry.last_accessed_at = datetime.now(UTC)
                await self.cache_entry_repository.save(cached_entry)
                await self.session.commit()
                return cached_path

        replicas = await self.blob_replica_repository.list_ready_replicas_with_backends(blob.id)
        for replica, backend in replicas:
            try:
                return await self.materialize_replica_to_local_cache(blob, replica, backend)
            except builtins.FileNotFoundError:
                continue
            except Exception:
                continue
        raise builtins.FileNotFoundError(blob.id)

    async def check_backend(self, backend: Backend) -> tuple[str, str | None]:
        adapter = self._build_adapter(backend)
        result = await adapter.check()
        return result["status"], result.get("detail")

    def _build_adapter(self, backend: Backend):
        config = json.loads(backend.config_json)
        if backend.type == "local":
            return LocalStorageAdapter(config["base_path"])
        if backend.type == "s3":
            return S3StorageAdapter(config)
        raise ValueError(f"Unsupported backend type: {backend.type}")

    @staticmethod
    def _build_storage_key(blob: Blob) -> str:
        content_hash = blob.content_hash
        prefix = content_hash[:2] if len(content_hash) >= 2 else "xx"
        return f"blobs/{prefix}/{content_hash}"