"""ImagesMixin — handler for image generation endpoints.

  POST /v1/images/generations   text-to-image, or image-to-image when
                                body includes an ``image`` (base64) field.
"""
from __future__ import annotations

import base64
import time


class ImagesMixin:
    """Provides image-generation endpoint handler.

    Relies on the concrete class for:
      self._resolve_ctx, self._send_json, self._not_supported, self._read_body.
    """

    def _handle_image_generations(self) -> None:
        body = self._read_body()
        if body is None:
            return

        ctx = self._resolve_ctx(body.get("model"))
        if ctx.infer_image is None:
            self._not_supported("/v1/images/generations")
            return

        prompt = body.get("prompt", "")
        if not prompt:
            self._send_json(400, {"error": "'prompt' is required"})
            return

        n    = int(body.get("n", 1))
        size = body.get("size") or None

        image_b64 = body.get("image")
        if image_b64:
            # ── image-to-image ────────────────────────────────────────────────
            if ctx.infer_image_edit is None:
                self._send_json(
                    501, {"error": "image-to-image is not supported by this model"},
                )
                return
            try:
                image_bytes = base64.b64decode(image_b64)
            except Exception:  # pylint: disable=broad-except
                self._send_json(400, {"error": "invalid base64 in 'image'"})
                return
            try:
                images = ctx.infer_image_edit(prompt, image_bytes, n, size)
            except Exception as exc:  # pylint: disable=broad-except
                self._send_json(500, {"error": str(exc)})
                return
        else:
            # ── text-to-image ─────────────────────────────────────────────────
            try:
                images = ctx.infer_image(prompt, n, size)
            except Exception as exc:  # pylint: disable=broad-except
                self._send_json(500, {"error": str(exc)})
                return

        self._send_json(200, {
            "created": int(time.time()),
            "data": [{"b64_json": img} for img in images],
        })
