OCR/docengine/app/storage/provider.py

from __future__ import annotations

import hashlib
import os
import shutil
import uuid
from abc import ABC, abstractmethod
from pathlib import Path

from app.core.config import settings
from app.core.exceptions import StorageError
from app.core.logging_config import get_logger

logger = get_logger(__name__)


class StorageProvider(ABC):
    """Abstract base class for storage providers."""

    @abstractmethod
    def save_file(self, file_data: bytes, directory: str, filename: str | None = None) -> str:
        """Save file data and return the storage path."""
        ...

    @abstractmethod
    def read_file(self, storage_path: str) -> bytes:
        """Read file data from storage."""
        ...

    @abstractmethod
    def delete_file(self, storage_path: str) -> bool:
        """Delete a file from storage. Returns True if successful."""
        ...

    @abstractmethod
    def file_exists(self, storage_path: str) -> bool:
        """Check if a file exists in storage."""
        ...

    @abstractmethod
    def get_file_size(self, storage_path: str) -> int:
        """Get the file size in bytes."""
        ...

    @abstractmethod
    def get_absolute_path(self, storage_path: str) -> str:
        """Get the absolute filesystem path for a storage path."""
        ...

    @staticmethod
    def compute_checksum(data: bytes, algorithm: str = "sha256") -> str:
        """Compute checksum of file data."""
        hasher = hashlib.new(algorithm)
        hasher.update(data)
        return hasher.hexdigest()

    @staticmethod
    def generate_filename(original_filename: str) -> str:
        """Generate a unique filename preserving the original extension."""
        ext = Path(original_filename).suffix.lower()
        return f"{uuid.uuid4().hex}{ext}"


class LocalStorageProvider(StorageProvider):
    """Local filesystem storage provider."""

    def __init__(self, base_path: str | None = None) -> None:
        self.base_path = Path(base_path or settings.storage_local_path).resolve()
        self._ensure_directories()

    def _ensure_directories(self) -> None:
        """Create required storage directories."""
        for subdir in ("documents", "templates", "images", "temp", "rendered"):
            (self.base_path / subdir).mkdir(parents=True, exist_ok=True)

    def _resolve_path(self, storage_path: str) -> Path:
        """Resolve a storage path to an absolute path."""
        resolved = (self.base_path / storage_path).resolve()
        if not str(resolved).startswith(str(self.base_path)):
            raise StorageError(f"Path traversal detected: {storage_path}")
        return resolved

    def save_file(self, file_data: bytes, directory: str, filename: str | None = None) -> str:
        """Save file data to local storage."""
        if filename is None:
            filename = f"{uuid.uuid4().hex}.bin"

        dir_path = self.base_path / directory
        dir_path.mkdir(parents=True, exist_ok=True)

        file_path = dir_path / filename
        try:
            file_path.write_bytes(file_data)
            storage_path = str(file_path.relative_to(self.base_path))
            logger.info("file_saved", storage_path=storage_path, size=len(file_data))
            return storage_path
        except OSError as e:
            raise StorageError(f"Failed to save file: {e}") from e

    def read_file(self, storage_path: str) -> bytes:
        """Read file data from local storage."""
        file_path = self._resolve_path(storage_path)
        if not file_path.exists():
            raise StorageError(f"File not found: {storage_path}")
        try:
            return file_path.read_bytes()
        except OSError as e:
            raise StorageError(f"Failed to read file: {e}") from e

    def delete_file(self, storage_path: str) -> bool:
        """Delete a file from local storage."""
        file_path = self._resolve_path(storage_path)
        if not file_path.exists():
            return False
        try:
            file_path.unlink()
            logger.info("file_deleted", storage_path=storage_path)
            return True
        except OSError as e:
            logger.error("file_delete_failed", storage_path=storage_path, error=str(e))
            raise StorageError(f"Failed to delete file: {e}") from e

    def file_exists(self, storage_path: str) -> bool:
        """Check if a file exists in local storage."""
        file_path = self._resolve_path(storage_path)
        return file_path.exists()

    def get_file_size(self, storage_path: str) -> int:
        """Get the file size in bytes."""
        file_path = self._resolve_path(storage_path)
        if not file_path.exists():
            raise StorageError(f"File not found: {storage_path}")
        return file_path.stat().st_size

    def get_absolute_path(self, storage_path: str) -> str:
        """Get the absolute filesystem path."""
        return str(self._resolve_path(storage_path))

    def save_temp_file(self, file_data: bytes, filename: str) -> str:
        """Save a temporary file."""
        return self.save_file(file_data, "temp", filename)

    def cleanup_temp(self) -> int:
        """Remove all files in the temp directory."""
        temp_dir = self.base_path / "temp"
        count = 0
        if temp_dir.exists():
            for item in temp_dir.iterdir():
                if item.is_file():
                    item.unlink()
                    count += 1
                elif item.is_dir():
                    shutil.rmtree(item)
                    count += 1
        logger.info("temp_cleanup", files_removed=count)
        return count


def get_storage_provider() -> StorageProvider:
    """Factory function to get the configured storage provider."""
    if settings.storage_provider == "local":
        return LocalStorageProvider()
    raise StorageError(f"Unknown storage provider: {settings.storage_provider}")