storage_manager.py

import os
import shutil
from pathlib import Path
from typing import Dict, Any, List

def resolve_storage_dir() -> Path:
    # Check if /data is mounted and writable (HF Spaces persistent storage bucket)
    candidate = Path("/data")
    if candidate.exists():
        try:
            test_file = candidate / ".write_test"
            test_file.touch(exist_ok=True)
            test_file.unlink(missing_ok=True)
            return candidate
        except Exception:
            pass
    # Local fallback
    local_data = Path("./data").resolve()
    local_data.mkdir(parents=True, exist_ok=True)
    return local_data

STORAGE_DIR = resolve_storage_dir()

# Ensure essential directories
RUNS_DIR = STORAGE_DIR / "runs"
ARTIFACTS_DIR = STORAGE_DIR / "artifacts"
MEMORY_DIR = STORAGE_DIR / "memory"

for d in (RUNS_DIR, ARTIFACTS_DIR, MEMORY_DIR):
    d.mkdir(parents=True, exist_ok=True)

def get_storage_info() -> Dict[str, Any]:
    try:
        total, used, free = shutil.disk_usage(STORAGE_DIR)
        is_hf_bucket = str(STORAGE_DIR).startswith("/data")
        file_count = sum(1 for _ in STORAGE_DIR.rglob("*") if _.is_file())
        return {
            "path": str(STORAGE_DIR),
            "is_hf_bucket": is_hf_bucket,
            "total_gb": round(total / (1024**3), 2),
            "used_gb": round(used / (1024**3), 2),
            "free_gb": round(free / (1024**3), 2),
            "percent_used": round((used / total) * 100, 1) if total > 0 else 0,
            "file_count": file_count,
        }
    except Exception as e:
        return {
            "path": str(STORAGE_DIR),
            "is_hf_bucket": False,
            "error": str(e),
            "total_gb": 0,
            "used_gb": 0,
            "free_gb": 0,
            "percent_used": 0,
            "file_count": 0,
        }

def list_storage_files(subpath: str = "") -> List[Dict[str, Any]]:
    target = (STORAGE_DIR / subpath.lstrip("/")).resolve()
    if not str(target).startswith(str(STORAGE_DIR)):
        raise ValueError("Path traversal denied.")
    if not target.exists():
        return []
    
    results = []
    for item in sorted(target.iterdir(), key=lambda p: (not p.is_dir(), p.name)):
        rel = item.relative_to(STORAGE_DIR).as_posix()
        results.append({
            "name": item.name,
            "path": rel,
            "is_dir": item.is_dir(),
            "size_bytes": item.stat().st_size if item.is_file() else 0,
            "modified_time": item.stat().st_mtime,
        })
    return results

def read_storage_file(rel_path: str) -> str:
    target = (STORAGE_DIR / rel_path.lstrip("/")).resolve()
    if not str(target).startswith(str(STORAGE_DIR)):
        raise ValueError("Path traversal denied.")
    if not target.is_file():
        raise FileNotFoundError(f"File not found: {rel_path}")
    return target.read_text(encoding="utf-8", errors="replace")

def write_storage_file(rel_path: str, content: str) -> str:
    target = (STORAGE_DIR / rel_path.lstrip("/")).resolve()
    if not str(target).startswith(str(STORAGE_DIR)):
        raise ValueError("Path traversal denied.")
    target.parent.mkdir(parents=True, exist_ok=True)
    target.write_text(content, encoding="utf-8")
    return str(target.relative_to(STORAGE_DIR))

def delete_storage_file(rel_path: str) -> bool:
    target = (STORAGE_DIR / rel_path.lstrip("/")).resolve()
    if not str(target).startswith(str(STORAGE_DIR)):
        raise ValueError("Path traversal denied.")
    if not target.exists():
        return False
    if target.is_dir():
        shutil.rmtree(target)
    else:
        target.unlink()
    return True