mosaic/utils/cache.py

"""Atomic download cache with hash verification for Mosaic."""
import os
import hashlib
import shutil
import logging
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)


class DownloadCache:
    """Manages atomic downloads with hash-sharded storage and integrity verification."""

    def __init__(self, base_dir: str = "output/raw"):
        self.base_dir = Path(base_dir)
        self.base_dir.mkdir(parents=True, exist_ok=True)

    def get_shard_path(self, content_hash: str, collection: str, ext: str = "") -> Path:
        """Get hash-sharded storage path: raw/{collection}/{ab}/{cd}/{hash}.{ext}"""
        if len(content_hash) < 4:
            raise ValueError(f"Invalid hash: {content_hash}")
        ab = content_hash[:2]
        cd = content_hash[2:4]
        filename = content_hash + (f".{ext}" if ext else "")
        return self.base_dir / collection / ab / cd / filename

    def store_atomic(self, tmp_path: str, content_hash: str, collection: str,
                     ext: str = "", expected_size: Optional[int] = None) -> str:
        """Atomically move a temp file to its cache location.

        Args:
            tmp_path: Path to the temporary download file
            content_hash: SHA256 hash of the content
            collection: Source collection name
            ext: File extension
            expected_size: Expected file size (from Content-Length)

        Returns:
            Final path of the cached file

        Raises:
            ValueError: If hash or size verification fails
        """
        tmp = Path(tmp_path)
        if not tmp.exists():
            raise FileNotFoundError(f"Temp file not found: {tmp_path}")

        # Verify size if expected
        if expected_size is not None:
            actual_size = tmp.stat().st_size
            if actual_size != expected_size:
                raise ValueError(
                    f"Size mismatch: expected {expected_size}, got {actual_size}"
                )

        # Verify hash
        actual_hash = self.compute_hash(tmp_path)
        if actual_hash != content_hash:
            raise ValueError(
                f"Hash mismatch: expected {content_hash}, got {actual_hash}"
            )

        # Atomic move to final location
        final_path = self.get_shard_path(content_hash, collection, ext)
        final_path.parent.mkdir(parents=True, exist_ok=True)

        shutil.move(str(tmp), str(final_path))
        logger.debug("Cached: %s → %s", tmp_path, final_path)
        return str(final_path)

    def is_cached(self, content_hash: str, collection: str, ext: str = "") -> bool:
        """Check if a file exists in cache with valid hash."""
        path = self.get_shard_path(content_hash, collection, ext)
        if not path.exists():
            return False
        actual_hash = self.compute_hash(str(path))
        if actual_hash != content_hash:
            logger.warning("Cache corruption detected: %s (expected %s, got %s)",
                          path, content_hash, actual_hash)
            return False
        return True

    def get_cached_path(self, content_hash: str, collection: str, ext: str = "") -> Optional[str]:
        """Get path to cached file, or None if not cached/corrupted."""
        if self.is_cached(content_hash, collection, ext):
            return str(self.get_shard_path(content_hash, collection, ext))
        return None

    def verify_file(self, file_path: str, expected_hash: str) -> bool:
        """Verify a file's hash matches expected."""
        actual = self.compute_hash(file_path)
        return actual == expected_hash

    def verify_all(self, collection: str) -> tuple[int, int, list[str]]:
        """Verify all files in a collection.

        Returns:
            (verified_count, mismatch_count, list_of_mismatched_paths)
        """
        collection_dir = self.base_dir / collection
        if not collection_dir.exists():
            return 0, 0, []

        verified = 0
        mismatched = 0
        bad_paths = []

        for root, _, files in os.walk(collection_dir):
            for fname in files:
                if fname.startswith('.') or fname.endswith('.tmp'):
                    continue
                fpath = os.path.join(root, fname)
                # Extract expected hash from filename
                expected_hash = Path(fname).stem
                actual_hash = self.compute_hash(fpath)
                if actual_hash == expected_hash:
                    verified += 1
                else:
                    mismatched += 1
                    bad_paths.append(fpath)

        return verified, mismatched, bad_paths

    @staticmethod
    def compute_hash(file_path: str) -> str:
        """Compute SHA256 hash of a file."""
        sha256 = hashlib.sha256()
        with open(file_path, 'rb') as f:
            for chunk in iter(lambda: f.read(8192), b''):
                sha256.update(chunk)
        return sha256.hexdigest()

    @staticmethod
    def hash_content(content: bytes) -> str:
        """Compute SHA256 hash of bytes content."""
        return hashlib.sha256(content).hexdigest()

    def check_disk_space(self, required_mb: int = 100) -> bool:
        """Check if enough disk space is available."""
        stat = shutil.disk_usage(str(self.base_dir))
        free_mb = stat.free // (1024 * 1024)
        if free_mb < required_mb:
            logger.warning("Low disk space: %d MB free, %d MB required",
                          free_mb, required_mb)
            return False
        return True