src/audio_video_generator/utils/files.py

"""File and directory utilities."""

import json
import os
import re
import shutil
import time
import uuid
import zipfile
from pathlib import Path
from typing import Any, Dict, List, Optional


def ensure_dir(path: str) -> str:
    """Ensure directory exists, creating if necessary.
    
    Args:
        path: Directory path
        
    Returns:
        Absolute path to directory
    """
    os.makedirs(path, exist_ok=True)
    return os.path.abspath(path)


def make_run_dir(base_dir: str) -> str:
    """Create a unique run directory with timestamp.
    
    Args:
        base_dir: Base directory for runs
        
    Returns:
        Path to newly created run directory
    """
    run_id = time.strftime("%Y%m%d_%H%M%S") + "_" + uuid.uuid4().hex[:8]
    run_dir = os.path.join(base_dir, run_id)
    ensure_dir(run_dir)
    return run_dir


def get_checkpoint_dir(run_dir: str) -> str:
    """Get or create checkpoint subdirectory.
    
    Args:
        run_dir: Run directory path
        
    Returns:
        Path to checkpoint directory
    """
    return ensure_dir(os.path.join(run_dir, "checkpoints"))


def write_text_file(path: str, content: Optional[str]) -> str:
    """Write text content to file.
    
    Args:
        path: File path
        content: Text content to write
        
    Returns:
        Path to written file
    """
    ensure_dir(os.path.dirname(path))
    with open(path, "w", encoding="utf-8") as f:
        f.write("" if content is None else str(content))
    return path


def write_json_file(path: str, data: Any) -> str:
    """Write JSON data to file.
    
    Args:
        path: File path
        data: Data to serialize as JSON
        
    Returns:
        Path to written file
    """
    ensure_dir(os.path.dirname(path))
    with open(path, "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
    return path


def extract_zip(zip_path: str, extract_dir: str) -> str:
    """Extract ZIP file to directory.
    
    Args:
        zip_path: Path to ZIP file
        extract_dir: Directory to extract to
        
    Returns:
        Path to extraction directory
    """
    ensure_dir(extract_dir)
    with zipfile.ZipFile(zip_path, "r") as zf:
        zf.extractall(extract_dir)
    return extract_dir


def collect_images_recursive(folder_path: str, extensions: tuple = (".png", ".jpg", ".jpeg")) -> List[str]:
    """Recursively collect image files from directory.
    
    Args:
        folder_path: Root directory to search
        extensions: Tuple of valid image extensions
        
    Returns:
        List of image file paths
    """
    image_paths: List[str] = []
    for root, _, files in os.walk(folder_path):
        for file in files:
            if file.lower().endswith(extensions):
                image_paths.append(os.path.join(root, file))
    return image_paths


def sort_images(image_paths: List[str]) -> List[str]:
    """Sort images by filename number, then alphabetically.
    
    Files starting with numbers are sorted numerically first.
    
    Args:
        image_paths: List of image paths
        
    Returns:
        Sorted list of image paths
    """
    def sort_key(path: str):
        name = os.path.basename(path)
        m = re.match(r"(\d+)", name)
        if m:
            return (0, int(m.group(1)), name.lower())
        return (1, name.lower())
    
    return sorted(image_paths, key=sort_key)


def copy_file_to_dir(src_path: str, dst_dir: str) -> str:
    """Copy file to destination directory.
    
    Args:
        src_path: Source file path
        dst_dir: Destination directory
        
    Returns:
        Path to copied file
    """
    ensure_dir(dst_dir)
    dst_path = os.path.join(dst_dir, os.path.basename(src_path))
    shutil.copy2(src_path, dst_path)
    return dst_path


def list_files_in_folder(folder_path: str) -> List[str]:
    """List all files in folder recursively.
    
    Args:
        folder_path: Directory to scan
        
    Returns:
        List of file paths
    """
    if not folder_path or not os.path.isdir(folder_path):
        return []
    
    files: List[str] = []
    for root, _, names in os.walk(folder_path):
        for name in names:
            full = os.path.join(root, name)
            if os.path.isfile(full):
                files.append(full)
    return sorted(files)


def find_first_file_by_ext(folder_path: str, extensions: tuple) -> Optional[str]:
    """Find first file matching extensions.
    
    Args:
        folder_path: Directory to search
        extensions: Tuple of extensions to match
        
    Returns:
        Path to first matching file or None
    """
    files = list_files_in_folder(folder_path)
    for f in files:
        if f.lower().endswith(extensions):
            return f
    return None


def find_all_images_in_folder(folder_path: str, extensions: tuple = (".png", ".jpg", ".jpeg")) -> List[str]:
    """Find all images in folder.
    
    Args:
        folder_path: Directory to search
        extensions: Image extensions to match
        
    Returns:
        List of image paths
    """
    if not folder_path or not os.path.isdir(folder_path):
        return []
    
    images: List[str] = []
    for root, _, names in os.walk(folder_path):
        for name in names:
            full = os.path.join(root, name)
            if os.path.isfile(full) and name.lower().endswith(extensions):
                images.append(full)
    return sort_images(images)