pig.py/storage.py

#!/usr/bin/env python3
"""
Image storage using content-addressable storage (MD5 hash).

Files are stored by MD5 hash, providing automatic deduplication.
Optionally uses filevault if available.
"""

import asyncio
import hashlib
from pathlib import Path
from typing import Optional

import aiofiles
import aiofiles.os

try:
    from filevault import FileVault
    HAS_FILEVAULT = True
except ImportError:
    HAS_FILEVAULT = False


class ImageVault:
    """
    Content-addressable image storage.

    Uses filevault if available, falls back to simple directory structure.
    Files are stored by MD5 hash for deduplication.

    Directory structure (fallback mode):
        vault/
        ├── 00/
        │   ├── 00abc123def456....jpg
        │   └── 00xyz789...png
        ├── 01/
        │   └── 01...
        └── ff/
            └── ff...
    """

    def __init__(self, vault_path: str = "vault"):
        self.vault_path = Path(vault_path)
        self.filevault: Optional[FileVault] = None
        self._initialized = False

    async def init(self) -> None:
        """Initialize the vault."""
        if self._initialized:
            return

        if HAS_FILEVAULT:
            self.filevault = await asyncio.to_thread(FileVault, str(self.vault_path))
        else:
            # Fallback: create directory structure
            await aiofiles.os.makedirs(self.vault_path, exist_ok=True)
            # Create subdirectories for first 2 chars of hash (256 buckets)
            for i in range(256):
                subdir = self.vault_path / f"{i:02x}"
                await aiofiles.os.makedirs(subdir, exist_ok=True)

        self._initialized = True

    async def store(self, md5_hash: str, data: bytes, extension: str = "") -> str:
        """
        Store data by MD5 hash.

        Args:
            md5_hash: MD5 hash of the data
            data: Raw file data
            extension: File extension (without dot)

        Returns:
            Storage path/key
        """
        if HAS_FILEVAULT and self.filevault:
            # filevault stores by content hash automatically
            return await asyncio.to_thread(self.filevault.put, data)
        else:
            # Fallback: store in subdirectory by first 2 chars
            subdir = md5_hash[:2]
            filename = f"{md5_hash}.{extension}" if extension else md5_hash
            filepath = self.vault_path / subdir / filename

            async with aiofiles.open(filepath, 'wb') as f:
                await f.write(data)

            return str(filepath)

    async def exists(self, md5_hash: str) -> bool:
        """Check if a file exists by MD5 hash."""
        if HAS_FILEVAULT and self.filevault:
            return await asyncio.to_thread(self.filevault.exists, md5_hash)
        else:
            # Check for any file starting with this hash
            subdir = self.vault_path / md5_hash[:2]

            def _check_exists():
                if not subdir.exists():
                    return False
                for f in subdir.iterdir():
                    if f.name.startswith(md5_hash):
                        return True
                return False

            return await asyncio.to_thread(_check_exists)

    async def get(self, md5_hash: str) -> Optional[bytes]:
        """Retrieve data by MD5 hash."""
        if HAS_FILEVAULT and self.filevault:
            return await asyncio.to_thread(self.filevault.get, md5_hash)
        else:
            subdir = self.vault_path / md5_hash[:2]

            def _find_file():
                if not subdir.exists():
                    return None
                for f in subdir.iterdir():
                    if f.name.startswith(md5_hash):
                        return f
                return None

            file_path = await asyncio.to_thread(_find_file)
            if file_path is None:
                return None

            async with aiofiles.open(file_path, 'rb') as file:
                return await file.read()

    async def get_path(self, md5_hash: str) -> Optional[Path]:
        """Get the filesystem path for a stored file."""
        if HAS_FILEVAULT and self.filevault:
            # filevault may not expose paths directly
            def _get_path():
                return self.filevault.path(md5_hash) if hasattr(self.filevault, 'path') else None

            path = await asyncio.to_thread(_get_path)
            return Path(path) if path else None
        else:
            subdir = self.vault_path / md5_hash[:2]

            def _find_file():
                if not subdir.exists():
                    return None
                for f in subdir.iterdir():
                    if f.name.startswith(md5_hash):
                        return f
                return None

            return await asyncio.to_thread(_find_file)

    async def delete(self, md5_hash: str) -> bool:
        """Delete a file by MD5 hash."""
        if HAS_FILEVAULT and self.filevault:
            return await asyncio.to_thread(self.filevault.delete, md5_hash)
        else:
            subdir = self.vault_path / md5_hash[:2]

            def _find_file():
                if not subdir.exists():
                    return None
                for f in subdir.iterdir():
                    if f.name.startswith(md5_hash):
                        return f
                return None

            file_path = await asyncio.to_thread(_find_file)
            if file_path is None:
                return False

            await aiofiles.os.remove(file_path)
            return True

    async def stats(self) -> dict:
        """Get vault statistics."""
        if HAS_FILEVAULT and self.filevault:
            def _get_stats():
                return {
                    'backend': 'filevault',
                    'count': len(self.filevault) if hasattr(self.filevault, '__len__') else -1
                }
            return await asyncio.to_thread(_get_stats)
        else:
            def _compute_stats():
                count = 0
                total_size = 0
                for subdir in self.vault_path.iterdir():
                    if subdir.is_dir():
                        for f in subdir.iterdir():
                            count += 1
                            total_size += f.stat().st_size
                return {
                    'backend': 'directory',
                    'count': count,
                    'total_size_bytes': total_size
                }

            return await asyncio.to_thread(_compute_stats)