ETPlot/gallery/api.py

"""
Main API for gallery generation.

Provides the primary entry point for programmatic gallery generation.
"""

import shutil
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, cast

from gallery.builder import build_gallery, copy_assets, get_template
from gallery.config import GalleryConfig, GallerySource


def generate(
    config: Optional[Union[GalleryConfig, str, Path]] = None,
    web_folder: Optional[Union[str, Path]] = None,
    sources: Optional[List[Union[GallerySource, Dict[str, Any]]]] = None,
    clean_first: bool = False,
    verbose: bool = False,
    source_to_update: Optional[GallerySource] = None,
) -> bool:
    """
    Generate a scientific gallery from plot sources.

    Can be called in two ways:
    1. With a GalleryConfig object
    2. With explicit parameters (web_folder and sources)

    Args:
        config: GalleryConfig object or path to YAML config file.
                If this is provided, other args are ignored.
        web_folder: Output directory for the gallery.
                   Required if config is not provided.
        sources: List of GallerySource objects or dicts.
                Required if config is not provided.
        clean_first: If True, removes and recreates the gallery directory.
                    If False (default), performs incremental update.
        verbose: If True, prints progress messages.
        source_to_update: Optional specific source to update. When provided,
                         only this source is regenerated (incremental mode).
                         Other sources in config are preserved in the index.
                         Only effective when clean_first is False.

    Returns:
        True if gallery generation was successful, False otherwise

    Raises:
        ValueError: If required arguments are missing or invalid
        TypeError: If config type is invalid

    Example:
        # Using GalleryConfig object
        from gallery import generate, GalleryConfig, GallerySource

        config = GalleryConfig(
            web_folder="/output/path",
            sources=[
                GallerySource(name="plots", path="/path/to/plots"),
            ]
        )
        success = generate(config, verbose=True)

        # Using explicit parameters
        success = generate(
            web_folder="/output/path",
            sources=[
                {"name": "plots", "path": "/path/to/plots"},
            ],
            verbose=True
        )

        # Loading from YAML config
        success = generate(config="config.yaml", verbose=True)
    """
    try:
        # Load or create configuration
        if config is not None:
            if isinstance(config, (str, Path)):
                config = GalleryConfig.from_yaml(config)
            elif not isinstance(config, GalleryConfig):
                raise TypeError(f"config must be GalleryConfig, str, or Path, got {type(config)}")
        else:
            if web_folder is None or sources is None:
                raise ValueError("Either config or both web_folder and sources must be provided")
            config = GalleryConfig(web_folder=web_folder, sources=sources or [])

        # Validate configuration
        if not config.sources:
            if verbose:
                print("Warning: No sources configured")
            return False

        # Check if web_folder is writable
        web_folder_path = Path(config.web_folder)
        if not _is_writable(web_folder_path):
            if verbose:
                print(f"Error: Cannot write to web_folder: {config.web_folder}")
            return False

        # Create gallery root directory
        gallery_root = web_folder_path / config.plot_root

        if clean_first and gallery_root.exists():
            if verbose:
                print(f"Cleaning gallery directory {gallery_root}...")
            try:
                shutil.rmtree(gallery_root)
            except Exception as e:
                if verbose:
                    print(f"Warning: Could not clean directory: {e}")
                return False
        elif source_to_update and gallery_root.exists():
            # Incremental mode: only clean the specific source subdirectory
            source_subdir = gallery_root / source_to_update.name
            if source_subdir.exists():
                if verbose:
                    print(f"Updating source directory {source_to_update.name}...")
                try:
                    shutil.rmtree(source_subdir)
                except Exception as e:
                    if verbose:
                        print(f"Warning: Could not clean source subdirectory {source_subdir}: {e}")
                    return False

        try:
            gallery_root.mkdir(parents=True, exist_ok=True)
        except Exception as e:
            if verbose:
                print(f"Error: Could not create gallery directory: {e}")
            return False

        # Copy assets
        if not copy_assets(config, verbose=verbose):
            if verbose:
                print("Warning: Could not copy assets")
            # Don't fail, continue with generation

        # Get template
        try:
            template = get_template()
        except Exception as e:
            if verbose:
                print(f"Error: Could not load template: {e}")
            return False

        # Process sources
        # config.sources is always List[GallerySource] after GalleryConfig.__post_init__ normalizes it.
        source_subdirs = []
        for source in config.sources:
            source = cast(GallerySource, source)
            # Skip sources not matching the update target (if specified)
            if source_to_update and source.name != source_to_update.name:
                # Still include them in the index if they exist
                source_web_dir = gallery_root / source.name
                if source_web_dir.exists():
                    source_subdirs.append(source.name)
                continue

            try:
                source_path = Path(source.path).resolve()

                # Validate source exists
                if not source_path.exists():
                    if verbose:
                        print(f"Warning: Source {source.path} does not exist. Skipping.")
                    continue

                source_web_dir = gallery_root / source.name
                try:
                    source_web_dir.mkdir(parents=True, exist_ok=True)
                except Exception as e:
                    if verbose:
                        print(f"Warning: Could not create directory {source_web_dir}: {e}")
                    continue

                source_subdirs.append(source.name)

                # Process source
                if source_path.is_file() and source_path.suffix == ".pdf":
                    # Single PDF file
                    from gallery.utils.processing import process_plot_files

                    item = process_plot_files(
                        config=config,
                        plot_file=source_path,
                        web_dir=source_web_dir,
                    )
                    from gallery.utils.processing import render_gallery_page

                    render_gallery_page(
                        config=config,
                        template=template,
                        web_dir=source_web_dir,
                        items=[item],
                        subdirs=[],
                        relative_path=Path(source.name),
                    )
                elif source_path.is_dir():
                    # Directory of plots
                    build_gallery(config, source_path, source_web_dir, template, Path(source.name))
                else:
                    if verbose:
                        print(f"Warning: Source {source.path} is neither a directory nor a PDF file. Skipping.")
                    continue

                if verbose:
                    print(f"Processed {source.name}: {source.path}")

            except Exception as e:
                if verbose:
                    print(f"Warning: Error processing source {source.name}: {e}")
                continue

        # Render gallery root index
        try:
            from gallery.utils.processing import render_gallery_page

            render_gallery_page(
                config=config,
                template=template,
                web_dir=gallery_root,
                items=[],
                subdirs=source_subdirs,
                relative_path=Path("."),
                title="Gallery Root",
            )
        except Exception as e:
            if verbose:
                print(f"Warning: Could not render gallery root: {e}")
            # Don't fail, gallery is still usable

        plot_count = _count_gallery_plots(gallery_root)
        if plot_count == 0:
            print(f"Warning: Gallery at {gallery_root} appears to be empty — no plot files found!")
        else:
            print(f"✓ Gallery generated at {gallery_root} — {plot_count} plots total")

        return True

    except Exception as e:
        if verbose:
            print(f"Error: Gallery generation failed: {e}")
        return False


def _count_gallery_plots(gallery_root: Path) -> int:
    """Recursively count plot files (PDFs and HTMLs, excluding index.html) in the gallery output."""
    count = 0
    for f in gallery_root.rglob("*"):
        if f.is_file() and f.suffix.lower() in (".pdf", ".html") and f.name != "index.html":
            count += 1
    return count


def _is_writable(path: Path) -> bool:
    """
    Check if a path is writable.

    Creates the directory if it doesn't exist.

    Args:
        path: Path to check

    Returns:
        True if writable, False otherwise
    """
    try:
        path.mkdir(parents=True, exist_ok=True)
        # Try to create a test file
        test_file = path / ".gallery_test"
        test_file.touch()
        test_file.unlink()
        return True
    except Exception:
        return False