"""Default real provider — Replicate (SDXL inpainting + IP-Adapter).

Sends the room image, the inpaint mask, the catalog reference image, optional
depth conditioning, and a prompt to a Replicate model that supports image-
conditioned inpainting, then returns the generated image.

Credentials and the model id come from the environment only (see config.py).
Input field names are configurable so different compatible models can be used
without code changes.
"""
from __future__ import annotations

import io

import requests
from PIL import Image

from ..config import settings
from .base import ImageGenProvider


def _to_png_file(img: Image.Image, name: str) -> io.BytesIO:
    buf = io.BytesIO()
    img.save(buf, format="PNG")
    buf.seek(0)
    buf.name = name  # helps the client infer the content type
    return buf


def _output_to_image(output) -> Image.Image:
    """Replicate may return a URL, a list of URLs, or FileOutput objects."""
    item = output[0] if isinstance(output, (list, tuple)) else output
    if hasattr(item, "read"):  # FileOutput (newer client)
        return Image.open(io.BytesIO(item.read())).convert("RGB")
    resp = requests.get(str(item), timeout=180)
    resp.raise_for_status()
    return Image.open(io.BytesIO(resp.content)).convert("RGB")


class ReplicateImageGenProvider(ImageGenProvider):
    name = "replicate"

    def __init__(self) -> None:
        if not settings.REPLICATE_API_TOKEN:
            raise RuntimeError("REPLICATE_API_TOKEN is not set")
        if not settings.REPLICATE_MODEL:
            raise RuntimeError("REPLICATE_MODEL is not set")
        import replicate  # imported lazily so mock mode needs no creds

        self._client = replicate.Client(api_token=settings.REPLICATE_API_TOKEN)

    def generate(
        self,
        room_image: Image.Image,
        mask: Image.Image,
        reference_image: Image.Image,
        depth_map: Image.Image | None = None,
        prompt: str = "",
        **kwargs,
    ) -> Image.Image:
        inputs = {
            settings.REPLICATE_INPUT_IMAGE_KEY: _to_png_file(room_image.convert("RGB"), "room.png"),
            settings.REPLICATE_INPUT_MASK_KEY: _to_png_file(mask.convert("L"), "mask.png"),
            settings.REPLICATE_INPUT_REFERENCE_KEY: _to_png_file(
                reference_image.convert("RGB"), "reference.png"
            ),
            settings.REPLICATE_INPUT_PROMPT_KEY: prompt,
        }
        if depth_map is not None and kwargs.get("send_depth", True):
            inputs.setdefault("depth_image", _to_png_file(depth_map.convert("RGB"), "depth.png"))
        # Allow callers/models to inject any additional inputs verbatim.
        inputs.update(kwargs.get("extra_inputs") or {})

        output = self._client.run(settings.REPLICATE_MODEL, input=inputs)
        return _output_to_image(output)