"""Serverless backend using Hugging Face Inference Providers.

Runs the ``dx8152/Qwen-Edit-2509-Multiple-angles`` model through the
``InferenceClient.image_to_image`` API. Requires an HF token but **no local
GPU**, so it works on a basic CPU Space or from a robot-arm client machine.
"""

from __future__ import annotations

import io
from typing import Optional

from PIL import Image

from ..config import ANGLES_INFERENCE_MODEL
from ..images import fit_image
from .base import ImageEditBackend


class InferenceProvidersBackend(ImageEditBackend):
    source = "hf_inference_providers"

    def __init__(
        self,
        token: str,
        image_size: int,
        provider: str = "auto",
        model: str = ANGLES_INFERENCE_MODEL,
    ) -> None:
        self.token = token
        self.image_size = image_size
        self.provider = provider
        self.model = model
        self._client = None

    def prepare(self) -> None:
        from huggingface_hub import InferenceClient  # local import: heavy dep

        if not self.token or not self.token.strip():
            raise RuntimeError("An HF token is required for the Inference Providers backend.")
        # The dx8152 multi-angle LoRA is currently served ONLY by WaveSpeed, so
        # provider="auto" fails to route it. Default to WaveSpeed unless the
        # caller explicitly pinned a different provider.
        provider = self.provider
        if provider in (None, "", "auto"):
            provider = "wavespeed"
        self._client = InferenceClient(provider=provider, api_key=self.token.strip())

    def edit(
        self,
        image: Image.Image,
        prompt: str,
        seed: int,
        num_inference_steps: int,
        true_guidance_scale: float,
    ) -> Image.Image:
        base = fit_image(image.convert("RGB"), self.image_size)
        if not prompt.strip():
            return base

        if self._client is None:
            raise RuntimeError("Backend not prepared; call prepare() first.")

        buffer = io.BytesIO()
        base.save(buffer, format="PNG")
        result = self._client.image_to_image(
            buffer.getvalue(),
            prompt=prompt,
            model=self.model,
        )
        if not isinstance(result, Image.Image):
            raise RuntimeError("Inference Providers returned an unexpected response type.")
        # Normalise back to the requested geometry.
        return fit_image(result.convert("RGB"), self.image_size)