OpenViking/openviking/pyagfs/helpers.py

"""Helper functions for common file operations in AGFS.

This module provides high-level helper functions for common operations:
- cp: Copy files/directories within AGFS
- upload: Upload files/directories from local filesystem to AGFS
- download: Download files/directories from AGFS to local filesystem
"""

from __future__ import annotations

from collections.abc import Iterator
from pathlib import Path
from typing import Any

from .async_client import ensure_same_encryption_account
from .protocols import AGFSByteStream, AGFSSyncClientProtocol


def _call_with_optional_ctx(
    method: Any, *args: Any, ctx: dict[str, str] | None = None, **kwargs: Any
) -> Any:
    """Call a sync AGFS method with ctx when the backend supports it."""
    if ctx is None:
        return method(*args, **kwargs)
    try:
        return method(*args, ctx=ctx, **kwargs)
    except TypeError as exc:
        if "unexpected keyword argument 'ctx'" not in str(exc):
            raise
        return method(*args, **kwargs)


def _try_copy_within_mount_fast_path(
    client: AGFSSyncClientProtocol,
    src: str,
    dst: str,
    *,
    fs_ctx: dict[str, str] | None = None,
) -> bool:
    """Attempt a same-mount verbatim copy and return whether the fast-path was used."""
    copy_within_mount = getattr(client, "copy_within_mount", None)
    if not callable(copy_within_mount):
        return False

    result = _call_with_optional_ctx(copy_within_mount, src, dst, ctx=fs_ctx)
    if isinstance(result, dict):
        return bool(result.get("performed", False))
    return bool(result)


def cp(
    client: AGFSSyncClientProtocol,
    src: str,
    dst: str,
    recursive: bool = False,
    stream: bool = False,
    fs_ctx: dict[str, str] | None = None,
) -> None:
    """Copy a file or directory within AGFS.

    Args:
        client: AGFS binding client instance
        src: Source path in AGFS
        dst: Destination path in AGFS
        recursive: If True, copy directories recursively
        stream: If True, use streaming for large files (memory efficient)

    Raises:
        AGFSClientError: If source doesn't exist or operation fails

    Examples:
        >>> cp(client, "/file.txt", "/backup/file.txt")  # Copy file
        >>> cp(client, "/dir", "/backup/dir", recursive=True)  # Copy directory
    """
    ensure_same_encryption_account(src, dst)

    # Check if source exists and get its type
    src_info = _call_with_optional_ctx(client.stat, src, ctx=fs_ctx)
    is_dir = src_info.get("isDir", False)

    if is_dir:
        if not recursive:
            raise ValueError(f"Cannot copy directory '{src}' without recursive=True")
        _copy_directory(client, src, dst, stream, fs_ctx=fs_ctx)
    else:
        _copy_file(client, src, dst, stream, fs_ctx=fs_ctx)


def upload(
    client: AGFSSyncClientProtocol,
    local_path: str,
    remote_path: str,
    recursive: bool = False,
    stream: bool = False,
) -> None:
    """Upload a file or directory from local filesystem to AGFS.

    Args:
        client: AGFS binding client instance
        local_path: Path to local file or directory
        remote_path: Destination path in AGFS
        recursive: If True, upload directories recursively
        stream: If True, use streaming for large files (memory efficient)

    Raises:
        FileNotFoundError: If local path doesn't exist
        AGFSClientError: If upload fails

    Examples:
        >>> upload(client, "/tmp/file.txt", "/remote/file.txt")  # Upload file
        >>> upload(client, "/tmp/data", "/remote/data", recursive=True)  # Upload directory
    """
    local = Path(local_path)

    if not local.exists():
        raise FileNotFoundError(f"Local path does not exist: {local_path}")

    if local.is_dir():
        if not recursive:
            raise ValueError(f"Cannot upload directory '{local_path}' without recursive=True")
        _upload_directory(client, local, remote_path, stream)
    else:
        _upload_file(client, local, remote_path, stream)


def download(
    client: AGFSSyncClientProtocol,
    remote_path: str,
    local_path: str,
    recursive: bool = False,
    stream: bool = False,
) -> None:
    """Download a file or directory from AGFS to local filesystem.

    Args:
        client: AGFS binding client instance
        remote_path: Path in AGFS
        local_path: Destination path on local filesystem
        recursive: If True, download directories recursively
        stream: If True, use streaming for large files (memory efficient)

    Raises:
        AGFSClientError: If remote path doesn't exist or download fails

    Examples:
        >>> download(client, "/remote/file.txt", "/tmp/file.txt")  # Download file
        >>> download(client, "/remote/data", "/tmp/data", recursive=True)  # Download directory
    """
    # Check if remote path exists and get its type
    remote_info = client.stat(remote_path)
    is_dir = remote_info.get("isDir", False)

    if is_dir:
        if not recursive:
            raise ValueError(f"Cannot download directory '{remote_path}' without recursive=True")
        _download_directory(client, remote_path, Path(local_path), stream)
    else:
        _download_file(client, remote_path, Path(local_path), stream)


# Internal helper functions


def _copy_file(
    client: AGFSSyncClientProtocol,
    src: str,
    dst: str,
    stream: bool,
    *,
    fs_ctx: dict[str, str] | None = None,
) -> None:
    """Copy a single file within AGFS.

    Copy is always expressed in logical AGFS operations so encryption remains
    transparent to the business layer.
    """
    # Ensure parent directory exists
    _ensure_remote_parent_dir(client, dst, fs_ctx=fs_ctx)

    if _try_copy_within_mount_fast_path(client, src, dst, fs_ctx=fs_ctx):
        return

    if stream:
        # The binding client returns bytes or an iterator of bytes, not an
        # HTTP response object with iter_content().
        _call_with_optional_ctx(
            client.write,
            dst,
            _iter_file_bytes(_call_with_optional_ctx(client.cat, src, stream=True, ctx=fs_ctx)),
            ctx=fs_ctx,
        )
    else:
        # Read entire file and write
        data = _call_with_optional_ctx(client.cat, src, ctx=fs_ctx)
        _call_with_optional_ctx(client.write, dst, data, ctx=fs_ctx)


def _copy_directory(
    client: AGFSSyncClientProtocol,
    src: str,
    dst: str,
    stream: bool,
    *,
    fs_ctx: dict[str, str] | None = None,
) -> None:
    """Recursively copy a directory within AGFS."""
    # Create destination directory
    try:
        _call_with_optional_ctx(client.mkdir, dst, ctx=fs_ctx)
    except Exception:
        # Directory might already exist, continue
        pass

    # List source directory contents
    items = _call_with_optional_ctx(client.ls, src, ctx=fs_ctx)

    for item in items:
        item_name = item["name"]
        src_path = f"{src.rstrip('/')}/{item_name}"
        dst_path = f"{dst.rstrip('/')}/{item_name}"

        if item.get("isDir", False):
            # Recursively copy subdirectory
            _copy_directory(client, src_path, dst_path, stream, fs_ctx=fs_ctx)
        else:
            # Copy file
            _copy_file(client, src_path, dst_path, stream, fs_ctx=fs_ctx)


def _upload_file(
    client: AGFSSyncClientProtocol,
    local_file: Path,
    remote_path: str,
    stream: bool,
) -> None:
    """Upload a single file to AGFS."""
    # Ensure parent directory exists in AGFS
    _ensure_remote_parent_dir(client, remote_path)

    if stream:
        client.write(remote_path, _iter_local_file_bytes(local_file))
    else:
        # Read entire file
        with open(local_file, "rb") as f:
            data = f.read()
        client.write(remote_path, data)


def _upload_directory(
    client: AGFSSyncClientProtocol,
    local_dir: Path,
    remote_path: str,
    stream: bool,
) -> None:
    """Recursively upload a directory to AGFS."""
    # Create remote directory
    try:
        client.mkdir(remote_path)
    except Exception:
        # Directory might already exist, continue
        pass

    # Walk through local directory
    for item in local_dir.iterdir():
        remote_item_path = f"{remote_path.rstrip('/')}/{item.name}"

        if item.is_dir():
            # Recursively upload subdirectory
            _upload_directory(client, item, remote_item_path, stream)
        else:
            # Upload file
            _upload_file(client, item, remote_item_path, stream)


def _download_file(
    client: AGFSSyncClientProtocol,
    remote_path: str,
    local_file: Path,
    stream: bool,
) -> None:
    """Download a single file from AGFS."""
    # Ensure parent directory exists locally
    local_file.parent.mkdir(parents=True, exist_ok=True)

    if stream:
        with open(local_file, "wb") as f:
            for chunk in _iter_file_bytes(client.cat(remote_path, stream=True)):
                f.write(chunk)
    else:
        # Read entire file
        data = client.cat(remote_path)
        with open(local_file, "wb") as f:
            f.write(data)


def _download_directory(
    client: AGFSSyncClientProtocol,
    remote_path: str,
    local_dir: Path,
    stream: bool,
) -> None:
    """Recursively download a directory from AGFS."""
    # Create local directory
    local_dir.mkdir(parents=True, exist_ok=True)

    # List remote directory contents
    items = client.ls(remote_path)

    for item in items:
        item_name = item["name"]
        remote_item_path = f"{remote_path.rstrip('/')}/{item_name}"
        local_item_path = local_dir / item_name

        if item.get("isDir", False):
            # Recursively download subdirectory
            _download_directory(client, remote_item_path, local_item_path, stream)
        else:
            # Download file
            _download_file(client, remote_item_path, local_item_path, stream)


def _ensure_remote_parent_dir(
    client: AGFSSyncClientProtocol,
    path: str,
    *,
    fs_ctx: dict[str, str] | None = None,
) -> None:
    """Ensure the parent directory exists for a remote path."""
    parent = "/".join(path.rstrip("/").split("/")[:-1])
    if parent and parent != "/":
        # Try to create parent directory (and its parents)
        _ensure_remote_dir_recursive(client, parent, fs_ctx=fs_ctx)


def _ensure_remote_dir_recursive(
    client: AGFSSyncClientProtocol,
    path: str,
    *,
    fs_ctx: dict[str, str] | None = None,
) -> None:
    """Recursively ensure a directory exists in AGFS."""
    if not path or path == "/":
        return

    # Check if directory already exists
    try:
        info = _call_with_optional_ctx(client.stat, path, ctx=fs_ctx)
        if info.get("isDir", False):
            return  # Directory exists
    except Exception:
        # Directory doesn't exist, need to create it
        pass

    # Ensure parent exists first
    parent = "/".join(path.rstrip("/").split("/")[:-1])
    if parent and parent != "/":
        _ensure_remote_dir_recursive(client, parent, fs_ctx=fs_ctx)

    # Create this directory
    try:
        _call_with_optional_ctx(client.mkdir, path, ctx=fs_ctx)
    except Exception:
        # Might already exist due to race condition, ignore
        pass


def _iter_file_bytes(data: bytes | AGFSByteStream) -> AGFSByteStream:
    """Normalize AGFS read results to a byte iterator."""
    if isinstance(data, bytes):
        return iter((data,))
    if isinstance(data, Iterator):
        return data
    return iter(data)


def _iter_local_file_bytes(local_file: Path, chunk_size: int = 8192) -> AGFSByteStream:
    """Yield local file content in chunks for streaming uploads."""
    with open(local_file, "rb") as file_obj:
        while True:
            chunk = file_obj.read(chunk_size)
            if not chunk:
                break
            yield chunk