code/hub_storage.py

"""Hugging Face Hub Cloud Storage Integration.

This module provides seamless integration with Hugging Face Hub for:
- Cloud-based file storage and retrieval
- Persistent workspace on HF Spaces
- Direct code writing to repositories
- Dataset-backed persistence
- **NEW: HF Storage Buckets support** (persistent cloud storage)

Architecture:
┌─────────────┐     ┌──────────────────┐     ┌─────────────────┐
│   AI Agent  │────▶│  hub_storage.py   │────▶│  HuggingFace Hub │
│             │     │  (this module)    │     │  (Cloud Storage) │
└─────────────┘     └──────────────────┘     └─────────────────┘
                           │              │
                           ▼              ▼
                  ┌──────────────────┐ ┌─────────────────┐
                  │  Local Fallback  │ │  HF Bucket       │
                  │  (when offline)  │ │  (Persistent)   │
                  └──────────────────┘ └─────────────────┘

Usage Modes:
1. LOCAL: Files stored in ./workspace/ (default)
2. HUB: Files synced to HF Hub repository
3. HYBRID: Local + Hub sync (recommended for Spaces)
4. BUCKET: Files stored in HF Storage Bucket (NEW - best for Spaces)

HF Bucket API Reference:
- create_bucket: Create a new storage bucket
- upload_file: Upload files to bucket
- list_bucket_tree: List bucket contents
- Buckets provide S3-like persistent storage on HF Hub

Example usage with buckets:
    >>> from code.hub_storage import init_hub_storage
    >>> storage = init_hub_storage(
    ...     token="hf_...",
    ...     repo_id="sonic-coder/sonicoder",
    ...     bucket_repo_id="sonic-coder/sonicoder-workspace-bucket",
    ...     mode="bucket"
    ... )
    >>> # Now AI can write directly to cloud storage!
    >>> storage.write_file("project/app.py", "print('Hello Cloud!')")
"""

from __future__ import annotations

import json
import logging
import os
import tempfile
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, Callable, Optional, Iterator
from contextlib import contextmanager

logger = logging.getLogger(__name__)


@dataclass
class HubConfig:
    """Configuration for HF Hub storage."""
    
    # Repository settings
    repo_id: str = ""  # e.g., "username/repo-name" or "username/workspace-data"
    repo_type: str = "space"  # "space", "model", or "dataset"
    
    # Authentication
    token: str = ""
    
    # Storage mode
    mode: str = "local"  # "local", "hub", "hybrid", or "bucket"
    
    # Workspace path (for local/hybrid)
    local_workspace: str = "./workspace"
    
    # Sync settings
    auto_sync: bool = True
    sync_interval: int = 60  # seconds
    
    # Dataset repo for persistent storage (recommended for Spaces)
    dataset_repo_id: str = ""  # e.g., "username/sonicoder-workspace"
    
    # HF Storage Bucket (NEW - persistent cloud storage)
    bucket_repo_id: str = ""  # e.g., "username/sonicoder-bucket"
    bucket_mount_path: str = "/data/bucket"  # Where to mount in Spaces

    def is_hub_enabled(self) -> bool:
        """Check if Hub storage is enabled."""
        return self.mode in ("hub", "hybrid", "bucket") and bool(self.repo_id) and bool(self.token)
    
    def is_dataset_storage_enabled(self) -> bool:
        """Check if dataset-based storage is enabled."""
        return bool(self.dataset_repo_id) and bool(self.token)
    
    def is_bucket_enabled(self) -> bool:
        """Check if HF Storage Bucket is enabled (NEW)."""
        return (
            self.mode == "bucket" and 
            bool(self.bucket_repo_id) and 
            bool(self.token)
        )


@dataclass 
class FileMetadata:
    """Metadata for a file in storage."""
    path: str
    size: int = 0
    modified_at: float = field(default_factory=time.time)
    created_at: float = field(default_factory=time.time)
    content_hash: str = ""
    is_local: bool = True
    is_remote: bool = False
    is_in_bucket: bool = False  # NEW: Track if file is in bucket


class HubStorageError(Exception):
    """Custom exception for Hub storage operations."""
    
    def __init__(self, message: str, operation: str = "", recoverable: bool = True):
        self.message = message
        self.operation = operation
        self.recoverable = recoverable
        super().__init__(f"[{operation}] {message}")


class HubStorage:
    """
    Main class for Hugging Face Hub cloud storage integration.
    
    Provides transparent file operations that work with both
    local filesystem, HF Hub repos, AND HF Storage Buckets.
    
    Example:
        >>> storage = HubStorage(
        ...     repo_id="sonic-coder/sonicoder",
        ...     token="hf_...",
        ...     mode="bucket",
        ...     bucket_repo_id="sonic-coder/sonicoder-workspace"
        ... )
        >>> storage.initialize()
        >>> storage.write_file("app.py", "print('Hello!')")  # Writes to cloud!
        >>> content = storage.read_file("app.py")
        >>> files = storage.list_files()
    """
    
    def __init__(
        self,
        config: Optional[HubConfig] = None,
        **kwargs
    ):
        self.config = config or HubConfig(**kwargs)
        self._api = None
        self._initialized = False
        self._local_workspace: Path = Path(self.config.local_workspace)
        self._sync_lock_needed = False
        
        # Cache for remote files
        self._remote_cache: dict[str, FileMetadata] = {}
        self._cache_timestamp: float = 0
        self._cache_ttl: int = 30  # seconds
        
        # Operation callbacks
        self._on_sync: Optional[Callable] = None
        self._on_error: Optional[Callable] = None
        
        # Bucket state
        self._bucket_created: bool = False
        self._bucket_mounted: bool = False
        
    @property
    def api(self):
        """Lazy initialization of HfApi."""
        if self._api is None:
            try:
                from huggingface_hub import HfApi
                self._api = HfApi(token=self.config.token)
            except ImportError:
                raise HubStorageError(
                    "huggingface_hub package not installed",
                    operation="initialize",
                    recoverable=False
                )
        return self._api
    
    def initialize(self) -> dict[str, Any]:
        """
        Initialize the Hub storage system.
        
        Creates necessary directories, validates connection,
        sets up the storage backend, and creates bucket if needed.
        
        Returns:
            Dict with initialization status and info.
        """
        result = {
            "success": False,
            "mode": self.config.mode,
            "hub_enabled": False,
            "dataset_enabled": False,
            "bucket_enabled": False,
            "workspace_path": str(self._local_workspace),
        }
        
        try:
            # Create local workspace directory
            self._local_workspace.mkdir(parents=True, exist_ok=True)
            
            # Validate Hub connection if enabled
            if self.config.is_hub_enabled():
                try:
                    # Test API connection by getting user info
                    user_info = self.api.whoami()
                    result["hub_enabled"] = True
                    result["user"] = user_info.get("name", "unknown")
                    logger.info("Connected to HF Hub as: %s", result["user"])
                    
                    # Ensure repo exists
                    self._ensure_repo_exists()
                    
                except Exception as e:
                    logger.warning("Hub connection failed: %s", e)
                    result["hub_error"] = str(e)
                    # Fall back to local mode
                    self.config.mode = "local"
            
            # Check dataset storage
            if self.config.is_dataset_storage_enabled():
                try:
                    self._ensure_dataset_repo_exists()
                    result["dataset_enabled"] = True
                    logger.info("Dataset storage ready: %s", self.config.dataset_repo_id)
                except Exception as e:
                    logger.warning("Dataset setup failed: %s", e)
                    result["dataset_error"] = str(e)
            
            # NEW: Initialize Bucket storage if enabled
            if self.config.is_bucket_enabled():
                try:
                    bucket_result = self._initialize_bucket()
                    result["bucket_enabled"] = bucket_result["success"]
                    result.update(bucket_result)
                    logger.info("Bucket storage ready: %s", self.config.bucket_repo_id)
                except Exception as e:
                    logger.warning("Bucket setup failed: %s", e)
                    result["bucket_error"] = str(e)
            
            self._initialized = True
            result["success"] = True
            
            logger.info(
                "HubStorage initialized: mode=%s, hub=%s, dataset=%s, bucket=%s",
                self.config.mode,
                result["hub_enabled"],
                result.get("dataset_enabled", False),
                result.get("bucket_enabled", False)
            )
            
        except Exception as e:
            logger.exception("Failed to initialize HubStorage")
            result["error"] = str(e)
        
        return result
    
    def _ensure_repo_exists(self) -> None:
        """Ensure the target repository exists."""
        try:
            self.api.repo_info(
                repo_id=self.config.repo_id,
                repo_type=self.config.repo_type
            )
        except Exception:
            # Repo doesn't exist, create it
            logger.info("Creating repo: %s", self.config.repo_id)
            self.api.create_repo(
                repo_id=self.config.repo_id,
                repo_type=self.config.repo_type,
                exist_ok=True,
                private=False
            )
    
    def _ensure_dataset_repo_exists(self) -> None:
        """Ensure the dataset repository exists for persistent storage."""
        try:
            self.api.repo_info(
                repo_id=self.config.dataset_repo_id,
                repo_type="dataset"
            )
        except Exception:
            logger.info("Creating dataset repo: %s", self.config.dataset_repo_id)
            self.api.create_repo(
                repo_id=self.config.dataset_repo_id,
                repo_type="dataset",
                exist_ok=True,
                private=True  # Workspaces are usually private
            )
    
    # ═══════════════════════════════════════════════════════════════════
    # HF STORAGE BUCKET METHODS (NEW)
    # ═══════════════════════════════════════════════════════════════════
    
    def _initialize_bucket(self) -> dict[str, Any]:
        """
        Initialize HF Storage Bucket.
        
        Creates the bucket if it doesn't exist and prepares it for use.
        Uses the new HuggingFace Hub Storage Bucket API.
        
        Returns:
            Dict with bucket initialization status.
        """
        result = {
            "success": False,
            "bucket_repo_id": self.config.bucket_repo_id,
            "created": False,
            "mounted": False,
        }
        
        try:
            # Try to get bucket info (check if exists)
            try:
                bucket_info = self.api.get_bucket_info(
                    repo_id=self.config.bucket_repo_id
                )
                result["exists"] = True
                logger.info("Bucket exists: %s", self.config.bucket_repo_id)
            except Exception:
                # Bucket doesn't exist, create it
                logger.info("Creating new bucket: %s", self.config.bucket_repo_id)
                
                # Use create_bucket API
                try:
                    create_result = self.api.create_bucket(
                        repo_id=self.config.bucket_repo_id,
                        description="SoniCoder AI Workspace - Persistent Cloud Storage",
                        public=False,  # Private workspace
                    )
                    result["created"] = True
                    self._bucket_created = True
                    logger.info("Bucket created successfully")
                except AttributeError:
                    # Fallback: create as dataset repo (older API)
                    logger.warning("create_bucket not available, using dataset fallback")
                    self.api.create_repo(
                        repo_id=self.config.bucket_repo_id,
                        repo_type="dataset",
                        exist_ok=True,
                        private=True
                    )
                    result["created"] = True
                    self._bucket_created = True
            
            # Check if running in Spaces with mount
            if os.environ.get("SPACE_ID"):
                mount_path = self.config.bucket_mount_path
                if os.path.exists(mount_path):
                    result["mounted"] = True
                    self._bucket_mounted = True
                    logger.info("Bucket mounted at: %s", mount_path)
            
            result["success"] = True
            
        except Exception as e:
            logger.exception("Failed to initialize bucket")
            result["error"] = str(e)
        
        return result
    
    def create_bucket(
        self,
        bucket_repo_id: str = "",
        description: str = "SoniCoder Workspace Bucket",
        public: bool = False
    ) -> dict[str, Any]:
        """
        Create a new HF Storage Bucket explicitly.
        
        This is the main method to create a bucket for cloud storage.
        The bucket provides persistent S3-like storage on HF Hub.
        
        Args:
            bucket_repo_id: Bucket ID (e.g., "username/my-bucket")
                          Uses config value if not provided
            description: Bucket description
            public: Whether bucket is publicly readable
            
        Returns:
            Dict with creation result.
            
        Example:
            >>> result = storage.create_bucket(
            ...     bucket_repo_id="sonic-coder/my-workspace",
            ...     description="AI Code Workspace"
            ... )
            >>> print(result)  # {"success": True, "bucket_url": "..."}
        """
        if not bucket_repo_id:
            bucket_repo_id = self.config.bucket_repo_id
        
        if not bucket_repo_id:
            return {
                "success": False,
                "error": "No bucket_repo_id provided"
            }
        
        result = {
            "success": False,
            "bucket_repo_id": bucket_repo_id,
        }
        
        try:
            # Try new bucket API first
            try:
                create_result = self.api.create_bucket(
                    repo_id=bucket_repo_id,
                    description=description,
                    public=public,
                )
                result["success"] = True
                result["created"] = True
                result["bucket_url"] = f"https://huggingface.co/{bucket_repo_id}"
                
                # Update config
                self.config.bucket_repo_id = bucket_repo_id
                self._bucket_created = True
                
                logger.info("Bucket created: %s", bucket_repo_id)
                
            except AttributeError:
                # Fallback: Use dataset as bucket-like storage
                logger.info("Using dataset as bucket storage")
                self.api.create_repo(
                    repo_id=bucket_repo_id,
                    repo_type="dataset",
                    exist_ok=True,
                    private=not public
                )
                result["success"] = True
                result["created"] = True
                result["fallback_mode"] = "dataset"
                result["bucket_url"] = f"https://huggingface.co/{bucket_repo_id}"
                
                self.config.bucket_repo_id = bucket_repo_id
                self._bucket_created = True
                
        except Exception as e:
            logger.exception("Failed to create bucket")
            result["error"] = str(e)
        
        return result
    
    def write_to_bucket(
        self,
        path: str,
        content: str | bytes,
        commit_message: str | None = None
    ) -> FileMetadata:
        """
        Write a file directly to HF Storage Bucket.
        
        This is the MAIN method for writing files to cloud storage.
        When using bucket mode, this writes directly to persistent
        cloud storage that survives Space restarts.
        
        Args:
            path: Relative file path (e.g., "project/main.py")
            content: File content (string or bytes)
            commit_message: Optional commit message
            
        Returns:
            FileMetadata with file information.
            
        Example:
            >>> # Write code directly to cloud!
            >>> meta = storage.write_to_bucket("hello.py", "print('Hello Cloud!')")
            >>> print(f"File saved to cloud: {meta.path}")
        """
        # Normalize path
        path = self._normalize_path(path)
        
        if not commit_message:
            commit_message = f"AI generated: {path}"
        
        # Determine target repo (bucket or fallback)
        target_repo = self.config.bucket_repo_id or self.config.dataset_repo_id or self.config.repo_id
        target_repo_type = "dataset" if self.config.bucket_repo_id or self.config.dataset_repo_id else self.config.repo_type
        
        # Write to local cache first
        local_path = self._local_workspace / path
        local_path.parent.mkdir(parents=True, exist_ok=True)
        
        if isinstance(content, bytes):
            local_path.write_bytes(content)
            content_size = len(content)
        else:
            local_path.write_text(content, encoding='utf-8')
            content_size = len(content.encode('utf-8'))
        
        # Upload to bucket/cloud
        metadata = FileMetadata(
            path=path,
            size=content_size,
            is_local=True,
            is_in_bucket=False
        )
        
        try:
            # Prepare content for upload
            if isinstance(content, str):
                content_bytes = content.encode('utf-8')
            else:
                content_bytes = content
            
            # For large files, use temp file
            if len(content_bytes) > 100000:  # > 100KB
                with tempfile.NamedTemporaryFile(
                    mode='wb',
                    suffix=Path(path).suffix,
                    delete=False
                ) as tmp:
                    tmp.write(content_bytes)
                    tmp_path = tmp.name
                
                try:
                    self.api.upload_file(
                        path_or_fileobj=tmp_path,
                        path_in_repo=path,
                        repo_id=target_repo,
                        repo_type=target_repo_type,
                        commit_message=commit_message
                    )
                finally:
                    if os.path.exists(tmp_path):
                        os.unlink(tmp_path)
            else:
                # Upload from memory
                from io import BytesIO
                self.api.upload_file(
                    path_or_fileobj=BytesIO(content_bytes),
                    path_in_repo=path,
                    repo_id=target_repo,
                    repo_type=target_repo_type,
                    commit_message=commit_message
                )
            
            # Update metadata
            metadata.is_remote = True
            metadata.is_in_bucket = bool(self.config.bucket_repo_id)
            
            logger.info("Written to bucket: %s (%d bytes)", path, content_size)
            
        except Exception as e:
            logger.error("Failed to write to bucket: %s", e)
            if self._on_error:
                self._on_error("write_to_bucket", path, e)
            raise HubStorageError(
                f"Failed to write to bucket: {e}",
                operation="write_to_bucket"
            )
        
        return metadata
    
    def read_from_bucket(self, path: str) -> str:
        """
        Read a file from HF Storage Bucket.
        
        Args:
            path: Relative file path in bucket
            
        Returns:
            File content as string.
            
        Raises:
            HubStorageError if file not found.
        """
        path = self._normalize_path(path)
        
        # Try local cache first
        local_path = self._local_workspace / path
        if local_path.exists():
            return local_path.read_text(encoding='utf-8')
        
        # Download from bucket
        target_repo = self.config.bucket_repo_id or self.config.dataset_repo_id or self.config.repo_id
        target_repo_type = "dataset" if self.config.bucket_repo_id or self.config.dataset_repo_id else self.config.repo_type
        
        try:
            from huggingface_hub import hf_hub_download
            
            local_path = hf_hub_download(
                repo_id=target_repo,
                filename=path,
                repo_type=target_repo_type,
                token=self.config.token
            )
            
            with open(local_path, 'r', encoding='utf-8') as f:
                return f.read()
                
        except Exception as e:
            raise HubStorageError(
                f"File not found in bucket: {path} - {e}",
                operation="read_from_bucket",
                recoverable=False
            )
    
    def list_bucket_files(
        self,
        path: str = ".",
        recursive: bool = True
    ) -> list[dict[str, Any]]:
        """
        List files in the HF Storage Bucket.
        
        Args:
            path: Directory path to list
            recursive: Whether to list recursively
            
        Returns:
            List of file info dicts.
        """
        target_repo = self.config.bucket_repo_id or self.config.dataset_repo_id or self.config.repo_id
        target_repo_type = "dataset" if self.config.bucket_repo_id or self.config.dataset_repo_id else self.config.repo_type
        
        files = []
        
        try:
            # Get repo info with file listing
            repo_info = self.api.repo_info(
                repo_id=target_repo,
                repo_type=target_repo_type
            )
            
            siblings = getattr(repo_info, 'siblings', [])
            
            for sibling in siblings:
                if hasattr(sibling, 'rfilename'):
                    file_path = sibling.rfilename
                    
                    # Filter by path prefix if specified
                    if path and path != ".":
                        if not file_path.startswith(path):
                            continue
                        if not recursive and '/' in file_path[len(path):]:
                            continue
                    
                    files.append({
                        "path": file_path,
                        "size": getattr(sibling, 'size', 0),
                        "blob_id": getattr(sibling, 'blob_id', ''),
                        "is_local": False,
                        "is_remote": True,
                        "is_in_bucket": True,
                    })
            
            logger.info("Listed %d files from bucket", len(files))
            
        except Exception as e:
            logger.error("Failed to list bucket files: %s", e)
        
        return files
    
    def delete_from_bucket(
        self,
        path: str,
        commit_message: str | None = None
    ) -> bool:
        """
        Delete a file from HF Storage Bucket.
        
        Args:
            path: File path to delete
            commit_message: Optional commit message
            
        Returns:
            True if deleted successfully.
        """
        path = self._normalize_path(path)
        
        if not commit_message:
            commit_message = f"Delete: {path}"
        
        target_repo = self.config.bucket_repo_id or self.config.dataset_repo_id or self.config.repo_id
        target_repo_type = "dataset" if self.config.bucket_repo_id or self.config.dataset_repo_id else self.config.repo_type
        
        try:
            self.api.delete_file(
                path_in_repo=path,
                repo_id=target_repo,
                repo_type=target_repo_type,
                commit_message=commit_message
            )
            
            # Delete from local cache too
            local_path = self._local_workspace / path
            if local_path.exists():
                local_path.unlink()
            
            logger.info("Deleted from bucket: %s", path)
            return True
            
        except Exception as e:
            logger.error("Failed to delete from bucket: %s", e)
            return False
    
    def sync_to_bucket(
        self,
        path: str | None = None,
        commit_message: str = "Sync to bucket"
    ) -> dict[str, Any]:
        """
        Sync local files to HF Bucket.
        
        Args:
            path: Specific file/directory to sync (None = all)
            commit_message: Commit message
            
        Returns:
            Sync result dict.
        """
        result = {
            "success": False,
            "uploaded": [],
            "skipped": [],
            "errors": []
        }
        
        if not self.config.bucket_repo_id and not self.config.dataset_repo_id:
            result["errors"].append("No bucket configured")
            return result
        
        try:
            files_to_sync = []
            
            if path:
                full_path = self._local_workspace / path
                if full_path.is_file():
                    files_to_sync.append(path)
                elif full_path.is_dir():
                    for fp in full_path.rglob("*"):
                        if fp.is_file():
                            files_to_sync.append(str(fp.relative_to(self._local_workspace)))
            else:
                for fp in self._local_workspace.rglob("*"):
                    if fp.is_file():
                        rel = str(fp.relative_to(self._local_workspace))
                        files_to_sync.append(rel)
            
            for file_path in files_to_sync:
                try:
                    local_path = self._local_workspace / file_path
                    content = local_path.read_text(encoding='utf-8')
                    
                    self.write_to_bucket(
                        file_path,
                        content,
                        f"{commit_message}: {file_path}"
                    )
                    
                    result["uploaded"].append(file_path)
                    
                except Exception as e:
                    logger.error("Failed to sync %s: %s", file_path, e)
                    result["errors"].append({"file": file_path, "error": str(e)})
            
            result["success"] = len(result["errors"]) == 0
            
        except Exception as e:
            result["errors"].append({"file": "*", "error": str(e)})
        
        return result
    
    def get_bucket_status(self) -> dict[str, Any]:
        """
        Get current bucket status and info.
        
        Returns:
            Dict with bucket status information.
        """
        status = {
            "available": False,
            "configured": False,
            "repo_id": self.config.bucket_repo_id,
            "mode": self.config.mode,
        }
        
        if not self.config.bucket_repo_id:
            return status
        
        status["configured"] = True
        
        try:
            # Try to get bucket info
            try:
                bucket_info = self.api.get_bucket_info(
                    repo_id=self.config.bucket_repo_id
                )
                status["available"] = True
                status["exists"] = True
                status.update({
                    k: v for k, v in bucket_info.items() 
                    if isinstance(v, (str, int, float, bool))
                })
            except AttributeError:
                # Fallback: check as dataset
                repo_info = self.api.repo_info(
                    repo_id=self.config.bucket_repo_id,
                    repo_type="dataset"
                )
                status["available"] = True
                status["exists"] = True
                status["fallback_mode"] = "dataset"
            
            # Count files
            files = self.list_bucket_files()
            status["file_count"] = len(files)
            status["total_size_mb"] = round(
                sum(f.get("size", 0) for f in files) / (1024 * 1024), 2
            )
            
        except Exception as e:
            status["error"] = str(e)
        
        return status
    
    # ═══════════════════════════════════════════════════════════════════
    # LEGACY FILE OPERATIONS (with bucket support)
    # ═══════════════════════════════════════════════════════════════════
    
    def write_file(
        self,
        path: str,
        content: str | bytes,
        commit_message: str | None = None,
        sync_to_hub: bool = True
    ) -> FileMetadata:
        """
        Write a file to storage.
        
        If bucket mode is enabled, writes directly to cloud bucket.
        Otherwise uses legacy behavior (local + optional hub sync).
        """
        # If bucket mode, use bucket method
        if self.config.is_bucket_enabled():
            return self.write_to_bucket(path, content, commit_message)
        
        # Legacy behavior
        path = self._normalize_path(path)
        
        if not commit_message:
            commit_message = f"Update {path}"
        
        # Write to local filesystem first
        local_path = self._local_workspace / path
        local_path.parent.mkdir(parents=True, exist_ok=True)
        
        if isinstance(content, bytes):
            local_path.write_bytes(content)
        else:
            local_path.write_text(content, encoding='utf-8')
        
        # Create metadata
        metadata = FileMetadata(
            path=path,
            size=local_path.stat().st_size,
            is_local=True
        )
        
        # Sync to Hub if enabled
        if sync_to_hub and self._should_sync_to_hub():
            try:
                metadata = self._upload_to_hub(path, content, commit_message)
                metadata.is_remote = True
            except Exception as e:
                logger.warning("Failed to sync %s to Hub: %s", path, e)
                if self._on_error:
                    self._on_error("write", path, e)
        
        return metadata
    
    def read_file(self, path: str, use_cache: bool = True) -> str:
        """
        Read a file from storage.
        
        If bucket mode is enabled, reads from bucket.
        Otherwise tries local then Hub.
        """
        # If bucket mode, use bucket method
        if self.config.is_bucket_enabled():
            return self.read_from_bucket(path)
        
        # Legacy behavior
        path = self._normalize_path(path)
        
        # Try local first
        local_path = self._local_workspace / path
        if local_path.exists():
            return local_path.read_text(encoding='utf-8')
        
        # Try Hub if enabled
        if self.config.is_hub_enabled() or self.config.is_dataset_storage_enabled():
            try:
                return self._download_from_hub(path)
            except Exception as e:
                logger.debug("Could not download %s from Hub: %s", path, e)
        
        raise HubStorageError(
            f"File not found: {path}",
            operation="read_file",
            recoverable=False
        )
    
    def read_binary(self, path: str) -> bytes:
        """Read a file as binary."""
        path = self._normalize_path(path)
        
        local_path = self._local_workspace / path
        if local_path.exists():
            return local_path.read_bytes()
        
        if self.config.is_hub_enabled() or self.config.is_bucket_enabled():
            with tempfile.NamedTemporaryFile(delete=False) as tmp:
                tmp_path = tmp.name
            
            try:
                self.api.hub_download(
                    repo_id=self._get_target_repo(),
                    filename=path,
                    repo_type=self._get_target_repo_type(),
                    local_dir=os.path.dirname(tmp_path),
                )
                
                downloaded = Path(tmp_path).parent / Path(path).name
                if downloaded.exists():
                    return downloaded.read_bytes()
            finally:
                if os.path.exists(tmp_path):
                    os.unlink(tmp_path)
        
        raise HubStorageError(f"Binary file not found: {path}", operation="read_binary")
    
    def delete_file(self, path: str, commit_message: str | None = None) -> bool:
        """Delete a file from storage."""
        path = self._normalize_path(path)
        
        # If bucket mode, use bucket delete
        if self.config.is_bucket_enabled():
            return self.delete_from_bucket(path, commit_message)
        
        if not commit_message:
            commit_message = f"Delete {path}"
        
        success = True
        
        # Delete from local
        local_path = self._local_workspace / path
        if local_path.exists():
            local_path.unlink()
            self._cleanup_empty_dirs(local_path.parent)
        
        # Delete from Hub
        if self._should_sync_to_hub():
            try:
                self.api.delete_file(
                    path_in_repo=path,
                    repo_id=self._get_target_repo(),
                    repo_type=self._get_target_repo_type(),
                    commit_message=commit_message
                )
            except Exception as e:
                logger.warning("Failed to delete %s from Hub: %s", path, e)
                success = False
        
        return success
    
    def file_exists(self, path: str) -> bool:
        """Check if a file exists in storage."""
        path = self._normalize_path(path)
        
        if (self._local_workspace / path).exists():
            return True
        
        if self.config.is_hub_enabled() or self.config.is_bucket_enabled():
            try:
                files = self.list_files(use_cache=True)
                return any(f["path"] == path for f in files)
            except Exception:
                pass
        
        return False
    
    def list_files(
        self,
        path: str = ".",
        recursive: bool = True,
        use_cache: bool = True
    ) -> list[dict[str, Any]]:
        """List files in storage."""
        # If bucket mode, use bucket listing
        if self.config.is_bucket_enabled():
            return self.list_bucket_files(path, recursive)
        
        files = []
        seen_paths = set()
        
        # List local files
        base_path = self._local_workspace / path
        if base_path.exists():
            for file_path in base_path.rglob("*") if recursive else base_path.iterdir():
                if file_path.is_file():
                    rel_path = str(file_path.relative_to(self._local_workspace))
                    stat = file_path.stat()
                    files.append({
                        "path": rel_path.replace(os.sep, "/"),
                        "size": stat.st_size,
                        "modified": stat.st_mtime,
                        "is_local": True,
                        "is_remote": False,
                    })
                    seen_paths.add(rel_path)
        
        # List remote files if hub enabled
        if self.config.is_hub_enabled() or self.config.is_dataset_storage_enabled():
            try:
                remote_files = self._list_hub_files(use_cache=use_cache)
                for rf in remote_files:
                    if rf["path"] not in seen_paths:
                        rf["is_local"] = False
                        rf["is_remote"] = True
                        files.append(rf)
            except Exception as e:
                logger.warning("Failed to list remote files: %e", e)
        
        files.sort(key=lambda x: x["path"])
        
        return files
    
    # ═══════════════════════════════════════════════════════════════════
    # LEGACY HUB OPERATIONS
    # ═══════════════════════════════════════════════════════════════════
    
    def _should_sync_to_hub(self) -> bool:
        """Check if we should sync to Hub based on config."""
        return (
            self.config.is_hub_enabled() or 
            self.config.is_dataset_storage_enabled()
        ) and self.config.auto_sync
    

    def _get_target_repo(self) -> str:
        """Get the target repository ID for operations."""
        if self.config.is_dataset_storage_enabled():
            return self.config.dataset_repo_id
        return self.config.repo_id
    
    def _get_target_repo_type(self) -> str:
        """Get the target repository type."""
        if self.config.is_dataset_storage_enabled():
            return "dataset"
        return self.config.repo_type
    
    def _upload_to_hub(
        self,
        path: str,
        content: str | bytes,
        commit_message: str
    ) -> FileMetadata:
        """Upload a file to HF Hub."""
        
        if isinstance(content, str) and len(content) > 100000:
            with tempfile.NamedTemporaryFile(
                mode='w',
                suffix=Path(path).suffix,
                delete=False,
                encoding='utf-8'
            ) as tmp:
                tmp.write(content)
                tmp_path = tmp.name
            
            try:
                self.api.upload_file(
                    path_or_fileobj=tmp_path,
                    path_in_repo=path,
                    repo_id=self._get_target_repo(),
                    repo_type=self._get_target_repo_type(),
                    commit_message=commit_message
                )
            finally:
                os.unlink(tmp_path)
        else:
            content_bytes = content.encode('utf-8') if isinstance(content, str) else content
            
            from io import BytesIO
            self.api.upload_file(
                path_or_fileobj=BytesIO(content_bytes),
                path_in_repo=path,
                repo_id=self._get_target_repo(),
                repo_type=self._get_target_repo_type(),
                commit_message=commit_message
            )
        
        logger.info("Uploaded %s to Hub (%s)", path, self._get_target_repo())
        
        return FileMetadata(
            path=path,
            size=len(content) if isinstance(content, str) else len(content_bytes),
            is_remote=True,
            is_local=(self._local_workspace / path).exists()
        )
    
    def _download_from_hub(self, path: str) -> str:
        """Download a file from HF Hub."""
        from huggingface_hub import hf_hub_download
        
        local_path = hf_hub_download(
            repo_id=self._get_target_repo(),
            filename=path,
            repo_type=self._get_target_repo_type(),
            token=self.config.token
        )
        
        with open(local_path, 'r', encoding='utf-8') as f:
            return f.read()
    
    def _list_hub_files(
        self,
        use_cache: bool = True
    ) -> list[dict[str, Any]]:
        """List files in the Hub repository."""
        
        if use_cache and self._remote_cache:
            if time.time() - self._cache_timestamp < self._cache_ttl:
                return list(self._remote_cache.values())
        
        repo_info = self.api.repo_info(
            repo_id=self._get_target_repo(),
            repo_type=self._get_target_repo_type()
        )
        
        files = []
        siblings = getattr(repo_info, 'siblings', [])
        
        for sibling in siblings:
            if hasattr(sibling, 'rfilename'):
                files.append({
                    "path": sibling.rfilename,
                    "size": getattr(sibling, 'size', 0),
                    "blob_id": getattr(sibling, 'blob_id', ''),
                    "is_local": False,
                    "is_remote": True,
                })
        
        self._remote_cache = {f["path"]: f for f in files}
        self._cache_timestamp = time.time()
        
        return files
    
    # ═══════════════════════════════════════════════════════════════════
    # SYNC OPERATIONS
    # ═══════════════════════════════════════════════════════════════════
    
    def sync_to_hub(
        self,
        path: str | None = None,
        commit_message: str = "Sync workspace"
    ) -> dict[str, Any]:
        """Sync local files to Hub/Bucket."""
        # If bucket mode, use bucket sync
        if self.config.is_bucket_enabled():
            return self.sync_to_bucket(path, commit_message)
        
        result = {
            "success": False,
            "uploaded": [],
            "skipped": [],
            "errors": []
        }
        
        if not self.config.is_hub_enabled() and not self.config.is_dataset_storage_enabled():
            result["errors"].append("Hub storage not configured")
            return result
        
        try:
            files_to_sync = []
            
            if path:
                full_path = self._local_workspace / path
                if full_path.is_file():
                    files_to_sync.append(path)
                elif full_path.is_dir():
                    for fp in full_path.rglob("*"):
                        if fp.is_file():
                            files_to_sync.append(str(fp.relative_to(self._local_workspace)))
            else:
                for fp in self._local_workspace.rglob("*"):
                    if fp.is_file():
                        rel = str(fp.relative_to(self._local_workspace))
                        files_to_sync.append(rel)
            
            for file_path in files_to_sync:
                try:
                    local_path = self._local_workspace / file_path
                    content = local_path.read_text(encoding='utf-8')
                    
                    self._upload_to_hub(
                        file_path,
                        content,
                        f"{commit_message}: {file_path}"
                    )
                    
                    result["uploaded"].append(file_path)
                    
                except Exception as e:
                    logger.error("Failed to sync %s: %s", file_path, e)
                    result["errors"].append({"file": file_path, "error": str(e)})
            
            result["success"] = len(result["errors"]) == 0
            
            if self._on_sync:
                self._on_sync(result)
                
        except Exception as e:
            result["errors"].append({"file": "*", "error": str(e)})
            logger.exception("Sync failed")
        
        return result
    
    def sync_from_hub(
        self,
        path: str | None = None
    ) -> dict[str, Any]:
        """Sync files from Hub to local."""
        result = {
            "success": False,
            "downloaded": [],
            "errors": []
        }
        
        try:
            if path:
                content = self._download_from_hub(path)
                local_path = self._local_workspace / path
                local_path.parent.mkdir(parents=True, exist_ok=True)
                local_path.write_text(content, encoding='utf-8')
                result["downloaded"].append(path)
            else:
                remote_files = self._list_hub_files(use_cache=False)
                
                for rf in remote_files:
                    try:
                        content = self._download_from_hub(rf["path"])
                        local_path = self._local_workspace / rf["path"]
                        local_path.parent.mkdir(parents=True, exist_ok=True)
                        local_path.write_text(content, encoding='utf-8')
                        result["downloaded"].append(rf["path"])
                    except Exception as e:
                        result["errors"].append({
                            "file": rf["path"],
                            "error": str(e)
                        })
            
            result["success"] = len(result["errors"]) == len(remote_files) == 0 or len(remote_files) == 0
            
        except Exception as e:
            result["errors"].append({"file": "*", "error": str(e)})
            logger.exception("Sync from hub failed")
        
        return result
    
    def full_sync(self) -> dict[str, Any]:
        """Perform bidirectional sync between local and Hub."""
        result = {
            "to_hub": {},
            "from_hub": {}
        }
        
        result["to_hub"] = self.sync_to_hub(commit_message="Full sync push")
        result["from_hub"] = self.sync_from_hub()
        
        return result
    
    # ═══════════════════════════════════════════════════════════════════
    # UTILITY METHODS
    # ═══════════════════════════════════════════════════════════════════
    
    def _normalize_path(self, path: str) -> str:
        """Normalize a file path."""
        if path.startswith("./"):
            path = path[2:]
        path = path.replace("\\", "/")
        while "//" in path:
            path = path.replace("//", "/")
        if ".." in path:
            raise HubStorageError(
                "Path traversal detected",
                operation="_normalize_path",
                recoverable=False
            )
        return path.lstrip("/")
    
    def _cleanup_empty_dirs(self, dir_path: Path) -> None:
        """Remove empty directories up to workspace root."""
        try:
            while dir_path != self._local_workspace:
                if dir_path.is_dir() and not any(dir_path.iterdir()):
                    dir_path.rmdir()
                    dir_path = dir_path.parent
                else:
                    break
        except Exception:
            pass
    
    def get_storage_stats(self) -> dict[str, Any]:
        """Get statistics about current storage."""
        stats = {
            "mode": self.config.mode,
            "hub_enabled": self.config.is_hub_enabled(),
            "dataset_enabled": self.config.is_dataset_storage_enabled(),
            "bucket_enabled": self.config.is_bucket_enabled(),
            "local_files": 0,
            "local_size_mb": 0,
            "last_sync": self._cache_timestamp,
        }
        
        # Count local files
        if self._local_workspace.exists():
            for f in self._local_workspace.rglob("*"):
                if f.is_file():
                    stats["local_files"] += 1
                    stats["local_size_mb"] += f.stat().st_size / (1024 * 1024)
        
        stats["local_size_mb"] = round(stats["local_size_mb"], 2)
        
        # Add bucket stats if available
        if self.config.is_bucket_enabled():
            try:
                bucket_status = self.get_bucket_status()
                stats["bucket_stats"] = bucket_status
            except Exception:
                pass
        
        return stats
    
    def clear_cache(self) -> None:
        """Clear the remote file cache."""
        self._remote_cache = {}
        self._cache_timestamp = 0
    
    def set_callbacks(
        self,
        on_sync: Callable | None = None,
        on_error: Callable | None = None
    ) -> None:
        """Set callback functions for events."""
        self._on_sync = on_sync
        self._on_error = on_error
    
    @contextmanager
    def batch_operation(self, commit_message: str = "Batch update"):
        """Context manager for batch operations."""
        original_auto_sync = self.config.auto_sync
        self.config.auto_sync = False
        
        files_written = []
        
        try:
            yield files_written
            
            if files_written:
                if self.config.is_bucket_enabled():
                    self.sync_to_bucket(commit_message=commit_message)
                else:
                    self.sync_to_hub(commit_message=commit_message)
        finally:
            self.config.auto_sync = original_auto_sync


# ════════════════════════════════════════════════════════════════════════
# GLOBAL INSTANCE MANAGEMENT
# ════════════════════════════════════════════════════════════════════════

_global_storage: Optional[HubStorage] = None


def get_hub_storage(**kwargs) -> HubStorage:
    """Get or create the global HubStorage instance."""
    global _global_storage
    if _global_storage is None:
        _global_storage = HubStorage(**kwargs)
    return _global_storage


def init_hub_storage(
    token: str = "",
    repo_id: str = "",
    dataset_repo_id: str = "",
    bucket_repo_id: str = "",  # NEW: Bucket support
    mode: str = "auto",
    **kwargs
) -> HubStorage:
    """
    Initialize the global HubStorage instance.
    
    Auto-detects environment when running on HF Spaces.
    Now supports HF Storage Buckets for persistent cloud storage!
    
    Args:
        token: HF API token
        repo_id: Main repository ID
        dataset_repo_id: Dataset repository for workspace storage
        bucket_repo_id: HF Storage Bucket ID (NEW - recommended for Spaces)
        mode: Storage mode ("local", "hub", "hybrid", "bucket", "auto")
        
    Returns:
        Initialized HubStorage instance.
        
    Example with bucket:
        >>> storage = init_hub_storage(
        ...     token="hf_...",
        ...     bucket_repo_id="sonic-coder/my-workspace-bucket",
        ...     mode="bucket"
        ... )
    """
    global _global_storage
    
    # Auto-detect mode
    if mode == "auto":
        if os.environ.get("SPACE_ID"):
            # On Spaces, prefer bucket if available
            if bucket_repo_id:
                mode = "bucket"
            elif token:
                mode = "hybrid"
            else:
                mode = "local"
        else:
            mode = "local"
    
    # Get token from env if not provided
    if not token:
        token = os.environ.get("HF_TOKEN", "") or os.environ.get("HF_AUTH_TOKEN", "")
    
    config = HubConfig(
        token=token,
        repo_id=repo_id,
        dataset_repo_id=dataset_repo_id,
        bucket_repo_id=bucket_repo_id,
        mode=mode,
        **kwargs
    )
    
    _global_storage = HubStorage(config=config)
    _global_storage.initialize()
    
    return _global_storage


def is_hub_available() -> bool:
    """Check if Hub storage is available and configured."""
    if _global_storage is None:
        return False
    return _global_storage.config.is_hub_enabled()


def is_bucket_available() -> bool:
    """Check if HF Bucket storage is available (NEW)."""
    if _global_storage is None:
        return False
    return _global_storage.config.is_bucket_enabled()


def get_storage_mode() -> str:
    """Get current storage mode."""
    if _global_storage is None:
        return "not_initialized"
    return _global_storage.config.mode