"""
I/O Utility Functions

Provides file system operations and tensor serialization utilities.
All operations are designed for safety and reproducibility.

Design Decisions:
- Atomic file operations where possible
- Type-safe tensor serialization
- Directory creation on demand
- No silent failures

Time Complexity: O(n) for file operations where n = file size
Space Complexity: O(n) for in-memory tensor operations
"""

from pathlib import Path
from typing import Any, Dict, Optional, Union
import numpy as np


def ensure_dir(path: Union[str, Path]) -> Path:
    """
    Ensure a directory exists, creating it if necessary.
    
    Args:
        path: Directory path (can be file path, will use parent)
        
    Returns:
        Path object for the directory
    """
    path = Path(path)
    
    # If path has an extension, assume it's a file and use parent
    if path.suffix:
        path = path.parent
    
    path.mkdir(parents=True, exist_ok=True)
    return path


def save_tensor(
    data: np.ndarray,
    path: Union[str, Path],
    metadata: Optional[Dict[str, Any]] = None
) -> None:
    """
    Save a NumPy tensor to disk with optional metadata.
    
    Uses NumPy's compressed format for storage efficiency.
    Metadata is stored alongside for reproducibility.
    
    Args:
        data: NumPy array to save
        path: Save path (will add .npz extension if not present)
        metadata: Optional metadata dictionary
        
    Example:
        >>> save_tensor(arr, "data/processed/train.npz", {"mean": 0.5})
    """
    path = Path(path)
    if path.suffix != ".npz":
        path = path.with_suffix(".npz")
    
    ensure_dir(path)
    
    save_dict = {"data": data}
    if metadata is not None:
        save_dict["metadata"] = np.array([metadata], dtype=object)
    
    np.savez_compressed(path, **save_dict)


def load_tensor(path: Union[str, Path]) -> tuple:
    """
    Load a NumPy tensor from disk with metadata.
    
    Args:
        path: Path to .npz file
        
    Returns:
        Tuple of (data array, metadata dict or None)
        
    Raises:
        FileNotFoundError: If file doesn't exist
    """
    path = Path(path)
    if path.suffix != ".npz":
        path = path.with_suffix(".npz")
    
    if not path.exists():
        raise FileNotFoundError(f"Tensor file not found: {path}")
    
    loaded = np.load(path, allow_pickle=True)
    data = loaded["data"]
    
    metadata = None
    if "metadata" in loaded:
        metadata = loaded["metadata"].item()
    
    return data, metadata


def get_file_size_mb(path: Union[str, Path]) -> float:
    """
    Get file size in megabytes.
    
    Args:
        path: Path to file
        
    Returns:
        File size in MB
    """
    path = Path(path)
    if not path.exists():
        return 0.0
    return path.stat().st_size / (1024 * 1024)


def list_files(
    directory: Union[str, Path],
    pattern: str = "*",
    recursive: bool = False
) -> list:
    """
    List files in a directory matching a pattern.
    
    Args:
        directory: Directory to search
        pattern: Glob pattern (e.g., "*.nc" for NetCDF files)
        recursive: If True, search subdirectories
        
    Returns:
        List of Path objects
    """
    directory = Path(directory)
    if not directory.exists():
        return []
    
    if recursive:
        return sorted(directory.rglob(pattern))
    return sorted(directory.glob(pattern))