src/backends/__init__.py

"""Image-edit backend selection with automatic fallback.

Priority:

1. **ZeroGPU Qwen** when running on a Hugging Face ZeroGPU Space (real
   multi-angle LoRA edits on an on-demand GPU).
2. **Local Qwen** when a CUDA GPU is present (free to run, fast 4-step).
3. **HF Inference Providers** serverless when an HF token is available
   (works with no local GPU).
4. **Geometric** approximation as a last resort (no GPU / no token).

This mirrors WakeForge's backend-selection pattern.
"""

from __future__ import annotations

from typing import Optional

from .base import ImageEditBackend
from .geometric import GeometricBackend
from .inference_providers import InferenceProvidersBackend
from .local_qwen import LocalQwenBackend
from .zerogpu import ZeroGpuQwenBackend, on_zerogpu

__all__ = [
    "ImageEditBackend",
    "GeometricBackend",
    "InferenceProvidersBackend",
    "LocalQwenBackend",
    "ZeroGpuQwenBackend",
    "select_backend",
]


def _cuda_available() -> bool:
    try:
        import torch  # local import: heavy dep

        return bool(torch.cuda.is_available())
    except Exception:  # noqa: BLE001 - torch missing or broken
        return False


def select_backend(
    hf_token: Optional[str],
    image_size: int,
    provider: str = "auto",
    prefer: str = "auto",
) -> ImageEditBackend:
    """Return a ready image-edit backend.

    ``prefer`` may be ``"auto"``, ``"local"``, ``"serverless"`` or
    ``"geometric"``. Never raises: falls back to a token-free geometric
    backend so the Space is always usable.
    """
    prefer = (prefer or "auto").lower()

    if prefer == "geometric":
        backend = GeometricBackend(image_size=image_size)
        backend.prepare()
        return backend

    # 1. ZeroGPU Space — the real Qwen multi-angle pipeline on an on-demand GPU.
    if prefer in ("auto", "local") and on_zerogpu():
        backend = ZeroGpuQwenBackend(image_size=image_size)
        try:
            backend.prepare()
            return backend
        except Exception as exc:  # noqa: BLE001 - fall through
            print(f"[backend] ZeroGPU Qwen unavailable ({exc}); trying next option.")

    # 2. Local CUDA GPU (dev machines / dedicated-GPU Spaces).
    want_local = prefer in ("auto", "local") and _cuda_available()
    if want_local:
        backend = LocalQwenBackend(image_size=image_size)
        try:
            backend.prepare()
            return backend
        except Exception as exc:  # noqa: BLE001 - fall through to serverless
            print(f"[backend] Local Qwen unavailable ({exc}); trying Inference Providers.")

    if prefer in ("auto", "serverless") and hf_token and hf_token.strip():
        try:
            backend = InferenceProvidersBackend(
                token=hf_token,
                image_size=image_size,
                provider=provider,
            )
            backend.prepare()
            return backend
        except Exception as exc:  # noqa: BLE001 - fall through to geometric
            print(f"[backend] Inference Providers unavailable ({exc}); using geometric fallback.")

    # Last resort: token-free, CPU-only geometric approximation.
    print("[backend] Using geometric fallback (no GPU / no HF token).")
    backend = GeometricBackend(image_size=image_size)
    backend.prepare()
    return backend