"""Background removal helpers for the ``prompt rembg`` command."""

from __future__ import annotations

from collections.abc import Callable
from importlib import import_module
from pathlib import Path
from typing import Any

SUPPORTED_REMBG_EXTENSIONS = frozenset({'.jpeg', '.jpg', '.png', '.webp'})

type SessionFactory = Callable[..., Any]
type RemoveFunction = Callable[..., bytes]


def run_rembg(
    input_path: Path,
    *,
    output_path: Path,
    model: str = 'u2net',
    force: bool = False,
) -> list[Path]:
    """Remove image backgrounds from a file or a directory.

    Args:
        input_path: Input image file path or directory containing images.
        output_path: Output file path for single image mode, or output directory
            for directory mode.
        model: rembg model name.
        force: Whether to overwrite existing output files.

    Returns:
        List of written output paths.

    Raises:
        ValueError: If paths are invalid or processing fails.
    """
    if not input_path.exists():
        raise ValueError(f'Input path not found: {input_path}')
    if input_path.is_file():
        return [remove_background(input_path, output_path=output_path, model=model, force=force)]
    if input_path.is_dir():
        return remove_background_batch(input_path, output_dir=output_path, model=model, force=force)
    raise ValueError(f'Input path must be a file or directory: {input_path}')


def remove_background(
    input_path: Path,
    *,
    output_path: Path,
    model: str = 'u2net',
    force: bool = False,
    session: Any | None = None,
    remove_fn: RemoveFunction | None = None,
) -> Path:
    """Remove background for a single image file.

    Args:
        input_path: Source image file.
        output_path: Destination image file.
        model: rembg model name.
        force: Whether to overwrite an existing output file.
        session: Optional pre-created rembg session for reuse.
        remove_fn: Optional preloaded rembg ``remove`` callable.

    Returns:
        Written output file path.

    Raises:
        ValueError: If validation, reading, or writing fails.
    """
    if not input_path.exists():
        raise ValueError(f'Input file not found: {input_path}')
    if not input_path.is_file():
        raise ValueError(f'Input path must be a file: {input_path}')
    _validate_image_path(input_path)
    _validate_file_output_path(output_path)
    if output_path.exists() and not force:
        raise ValueError(f'Output file already exists: {output_path}. Use --force to overwrite.')

    if session is None or remove_fn is None:
        new_session, loaded_remove = _load_rembg_api()
        if session is None:
            session = new_session(model_name=model)
        if remove_fn is None:
            remove_fn = loaded_remove

    parent = output_path.parent
    if parent.exists() and not parent.is_dir():
        raise ValueError(f'Output parent path must be a directory: {parent}')
    parent.mkdir(parents=True, exist_ok=True)

    try:
        input_bytes = input_path.read_bytes()
    except OSError as exc:
        raise ValueError(f'Failed to read input file {input_path}: {exc}') from exc

    try:
        output_bytes = remove_fn(input_bytes, session=session)
    except Exception as exc:
        raise ValueError(f'Failed to remove background for {input_path}: {exc}') from exc

    try:
        output_path.write_bytes(output_bytes)
    except OSError as exc:
        raise ValueError(f'Failed to write output file {output_path}: {exc}') from exc
    return output_path


def remove_background_batch(
    input_dir: Path,
    *,
    output_dir: Path,
    model: str = 'u2net',
    force: bool = False,
) -> list[Path]:
    """Remove backgrounds for every supported image in a directory.

    Only top-level files are processed; nested directories are ignored.

    Args:
        input_dir: Directory containing source images.
        output_dir: Directory where PNG files will be written.
        model: rembg model name.
        force: Whether to overwrite existing output files.

    Returns:
        Written output file paths.

    Raises:
        ValueError: If validation or processing fails.
    """
    if not input_dir.exists():
        raise ValueError(f'Input directory not found: {input_dir}')
    if not input_dir.is_dir():
        raise ValueError(f'Input path must be a directory: {input_dir}')
    _validate_directory_output_path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    image_files = _collect_image_files(input_dir)
    new_session, remove_fn = _load_rembg_api()
    session = new_session(model_name=model)

    written_files: list[Path] = []
    for image_file in image_files:
        destination = output_dir / f'{image_file.stem}.png'
        written_files.append(
            remove_background(
                image_file,
                output_path=destination,
                model=model,
                force=force,
                session=session,
                remove_fn=remove_fn,
            )
        )
    return written_files


def _collect_image_files(input_dir: Path) -> list[Path]:
    """Collect supported top-level image files from a directory."""
    files = sorted(
        (path for path in input_dir.iterdir() if path.is_file() and path.suffix.lower() in SUPPORTED_REMBG_EXTENSIONS),
        key=lambda path: path.name,
    )
    if not files:
        supported = ', '.join(sorted(SUPPORTED_REMBG_EXTENSIONS))
        raise ValueError(f'No supported image files found in directory: {input_dir}. Supported extensions: {supported}')
    return files


def _validate_image_path(path: Path) -> None:
    """Validate a single image input path."""
    if path.suffix.lower() not in SUPPORTED_REMBG_EXTENSIONS:
        supported = ', '.join(sorted(SUPPORTED_REMBG_EXTENSIONS))
        raise ValueError(f'Unsupported image format: {path.name}. Supported extensions: {supported}')


def _validate_file_output_path(output_path: Path) -> None:
    """Validate that output path is usable for single-file mode."""
    if output_path.exists() and output_path.is_dir():
        raise ValueError(f'Output path must be a file: {output_path}')


def _validate_directory_output_path(output_dir: Path) -> None:
    """Validate that output path is usable for directory mode."""
    if output_dir.exists() and not output_dir.is_dir():
        raise ValueError(f'Output path must be a directory: {output_dir}')


def _load_rembg_api() -> tuple[SessionFactory, RemoveFunction]:
    """Load rembg APIs lazily so tests do not require the dependency."""
    try:
        rembg_module = import_module('rembg')
    except ImportError as exc:
        raise ValueError('rembg CPU dependencies are not installed. Install the package with rembg[cpu].') from exc
    return rembg_module.new_session, rembg_module.remove
