"""File storage management for persisting UploadFile inputs."""

import os
import shutil
import uuid
from pathlib import Path
from typing import Any, Dict, BinaryIO
from datetime import datetime

# Try to import UploadFile from starlette/fastapi, but don't fail if missing
try:
    from starlette.datastructures import UploadFile
except ImportError:
    # Mock for type checking if not available
    class UploadFile:
        filename: str
        content_type: str
        file: BinaryIO
        size: int | None


class StorageManager:
    """Manages file persistence for workflow inputs."""

    def __init__(self, base_path: str = ".ombra_storage"):
        self.base_path = Path(base_path)
        self._init_storage()

    def _init_storage(self):
        """Initialize storage directory."""
        if not self.base_path.exists():
            self.base_path.mkdir(parents=True, exist_ok=True)

    def save_file(self, file: UploadFile) -> Dict[str, Any]:
        """Save an UploadFile to disk and return metadata.

        Args:
            file: The UploadFile object to save

        Returns:
            Metadata dictionary including the path to the saved file
        """
        # Create a unique path: YYYY/MM/UUID_filename
        now = datetime.now()
        relative_dir = Path(f"{now.year}") / f"{now.month:02d}"
        storage_dir = self.base_path / relative_dir
        
        if not storage_dir.exists():
            storage_dir.mkdir(parents=True, exist_ok=True)

        # Generate safe filename
        safe_filename = file.filename or "unnamed_file"
        unique_id = str(uuid.uuid4())[:8]
        filename = f"{unique_id}_{safe_filename}"
        file_path = storage_dir / filename

        # Save content
        try:
            # Ensure we are at the start of the file
            file.file.seek(0)
            
            with open(file_path, "wb") as dest:
                shutil.copyfileobj(file.file, dest)
                
            # Reset file pointer for subsequent usage in the current execution
            file.file.seek(0)
            
            # Get actual size
            size = file_path.stat().st_size

            return {
                "__type__": "OmbraFile",
                "filename": safe_filename,
                "content_type": file.content_type,
                "path": str(file_path),
                "size": size,
                "saved_at": now.isoformat()
            }
        except Exception as e:
            print(f"❌ Failed to persist file: {e}")
            # Return a fallback dict so we don't crash, but resume won't work for this file
            return {
                "__type__": "UploadFile", # Standard non-persisted type
                "filename": file.filename,
                "content_type": file.content_type,
                "size": file.size,
                "error": f"Persistence failed: {str(e)}"
            }

    def load_file(self, metadata: Dict[str, Any]) -> Any:
        """Reconstruct a file-like object from metadata.

        Args:
            metadata: The metadata dictionary saved by save_file

        Returns:
            A RehydratedFile object that mimics UploadFile
        """
        if metadata.get("__type__") != "OmbraFile":
            return None

        file_path = Path(metadata["path"])
        if not file_path.exists():
            raise FileNotFoundError(f"Stored file not found: {file_path}")

        return RehydratedFile(
            path=file_path,
            filename=metadata["filename"],
            content_type=metadata.get("content_type"),
            size=metadata.get("size")
        )


class RehydratedFile:
    """Mimics Starlette's UploadFile but backed by a stored file on disk."""

    def __init__(self, path: Path, filename: str, content_type: str | None = None, size: int | None = None):
        self.path = path
        self.filename = filename
        self.content_type = content_type
        self._size = size
        self._file = open(path, "rb")

    @property
    def file(self):
        return self._file

    @property
    def size(self):
        return self._size

    async def read(self, size: int = -1) -> bytes:
        return self._file.read(size)

    async def seek(self, offset: int) -> None:
        self._file.seek(offset)

    async def close(self) -> None:
        self._file.close()
        
    def to_metadata(self) -> Dict[str, Any]:
        """Return the metadata dictionary for persistence."""
        return {
            "__type__": "OmbraFile",
            "filename": self.filename,
            "content_type": self.content_type,
            "path": str(self.path),
            "size": self._size
        }
        
    def __repr__(self):
        return f"<RehydratedFile filename={self.filename} path={self.path}>"


# Global instance
_storage_manager: StorageManager | None = None

def get_storage_manager() -> StorageManager:
    global _storage_manager
    if _storage_manager is None:
        _storage_manager = StorageManager()
    return _storage_manager
