ifixkart-backend/app/services/migration_engine/storage_manager.py

import os
import uuid
import hashlib
import struct
from pathlib import Path
from typing import Tuple, Dict, Any, Optional

from app.models.MigrationModel import MediaItemStatusEnum

project_root = Path(__file__).resolve().parents[4]
MIGRATION_BASE_DIR = str(project_root / "uploads" / "migrations")
os.makedirs(MIGRATION_BASE_DIR, exist_ok=True)

class StorageManager:
    """
    Handles streaming disk storage, deterministic path resolution (/media/{sha256[0:2]}/{sha256}.ext),
    and atomic POSIX file writes (.tmp_xyz -> fsync -> rename) to guarantee zero corrupt files.
    """

    @staticmethod
    def get_job_dir(job_id: str) -> str:
        job_dir = os.path.join(MIGRATION_BASE_DIR, job_id)
        os.makedirs(os.path.join(job_dir, "dataset"), exist_ok=True)
        os.makedirs(os.path.join(job_dir, "archives"), exist_ok=True)
        os.makedirs(os.path.join(job_dir, "media"), exist_ok=True)
        os.makedirs(os.path.join(job_dir, "logs"), exist_ok=True)
        return job_dir

    @classmethod
    def save_job_config(cls, job_id: str, config: Dict[str, Any]):
        import json
        job_dir = cls.get_job_dir(job_id)
        config_path = os.path.join(job_dir, "job_config.json")
        with open(config_path, "w", encoding="utf-8") as f:
            json.dump(config, f)

    @classmethod
    def get_job_config(cls, job_id: str) -> Dict[str, Any]:
        import json
        job_dir = cls.get_job_dir(job_id)
        config_path = os.path.join(job_dir, "job_config.json")
        if os.path.exists(config_path):
            try:
                with open(config_path, "r", encoding="utf-8") as f:
                    return json.load(f)
            except Exception:
                pass
        return {}


    @classmethod
    async def save_upload_stream_async(cls, upload_file, dest_path: str, chunk_size: int = 256 * 1024) -> Tuple[int, str]:
        """
        Asynchronously streams HTTP upload file to disk using threadpool offloading and 256KB buffer.
        Prevents blocking asyncio event loop during multi-hundred MB uploads.
        """
        import asyncio

        def _write_sync():
            sha256_hash = hashlib.sha256()
            total_size = 0

            os.makedirs(os.path.dirname(dest_path), exist_ok=True)
            temp_path = f"{dest_path}.tmp_{uuid.uuid4().hex}"

            with open(temp_path, "wb") as f:
                while True:
                    chunk = upload_file.file.read(chunk_size)
                    if not chunk:
                        break
                    f.write(chunk)
                    sha256_hash.update(chunk)
                    total_size += len(chunk)
                f.flush()
                os.fsync(f.fileno())

            os.rename(temp_path, dest_path)
            return total_size, sha256_hash.hexdigest()

        return await asyncio.to_thread(_write_sync)

    @classmethod
    def save_upload_stream_sync(cls, file_obj, dest_path: str, chunk_size: int = 65536) -> Tuple[int, str]:
        """
        Synchronous 64KB streaming file saver.
        """
        sha256_hash = hashlib.sha256()
        total_size = 0

        os.makedirs(os.path.dirname(dest_path), exist_ok=True)
        temp_path = f"{dest_path}.tmp_{uuid.uuid4().hex}"

        with open(temp_path, "wb") as f:
            while chunk := file_obj.read(chunk_size):
                f.write(chunk)
                sha256_hash.update(chunk)
                total_size += len(chunk)
            f.flush()
            os.fsync(f.fileno())

        os.rename(temp_path, dest_path)
        return total_size, sha256_hash.hexdigest()

    @classmethod
    def write_media_file_atomically(cls, job_id: str, file_bytes: bytes, original_filename: str, sha256_hash: str) -> Tuple[str, str, int, int]:
        """
        Writes image file atomically: write to .tmp_xyz -> fsync -> rename to /media/{sha256[0:2]}/{sha256}.ext
        Returns: (storage_path, cdn_url, width, height)
        """
        ext = os.path.splitext(original_filename)[1].lower() or ".jpg"
        if not ext.startswith("."):
            ext = f".{ext}"

        sub_dir = sha256_hash[:2]
        media_dir = os.path.join(cls.get_job_dir(job_id), "media", sub_dir)
        os.makedirs(media_dir, exist_ok=True)

        filename = f"{sha256_hash}{ext}"
        final_disk_path = os.path.join(media_dir, filename)
        temp_disk_path = os.path.join(media_dir, f".tmp_{uuid.uuid4().hex}")

        # Fast binary PNG width/height header inspection
        width, height = 0, 0
        if ext == ".png" and len(file_bytes) >= 24:
            if file_bytes[:8] == b'\x89PNG\r\n\x1a\n' and file_bytes[12:16] == b'IHDR':
                try:
                    width, height = struct.unpack(">II", file_bytes[16:24])
                except Exception:
                    pass

        if width == 0 or height == 0:
            try:
                from PIL import Image as PIL_Image
                import io
                with PIL_Image.open(io.BytesIO(file_bytes)) as img:
                    width, height = img.size
            except Exception:
                width, height = 0, 0

        # Write to temporary file -> fsync -> rename
        if not os.path.exists(final_disk_path):
            with open(temp_disk_path, "wb") as f:
                f.write(file_bytes)
                f.flush()
                os.fsync(f.fileno())
            os.rename(temp_disk_path, final_disk_path)

        storage_path = os.path.join("uploads", "migrations", job_id, "media", sub_dir, filename)
        cdn_url = f"/uploads/migrations/{job_id}/media/{sub_dir}/{filename}"

        return storage_path, cdn_url, width, height

    @classmethod
    def exists(cls, storage_path: str) -> bool:
        """
        Verifies whether a physical storage file exists on disk.
        """
        if not storage_path:
            return False
        if storage_path.startswith("/"):
            return os.path.exists(storage_path)
        disk_path = str(project_root / storage_path)
        return os.path.exists(disk_path)