"""Async image generation and editing clients."""

from __future__ import annotations

import logging
from typing import TYPE_CHECKING

import openai

if TYPE_CHECKING:
    from io import BufferedReader, BytesIO

logger = logging.getLogger(__name__)

_image_gen_client: openai.AsyncOpenAI | None = None
_image_edit_client: openai.AsyncOpenAI | None = None


def get_image_gen_client() -> openai.AsyncOpenAI:
    """Return the shared async image-generation client, building it once."""
    global _image_gen_client
    if _image_gen_client is None:
        from vibe_bot.config import IMAGE_GEN_ENDPOINT, IMAGE_GEN_ENDPOINT_KEY

        _image_gen_client = openai.AsyncOpenAI(
            base_url=IMAGE_GEN_ENDPOINT,
            api_key=IMAGE_GEN_ENDPOINT_KEY,
            max_retries=0,
        )
    return _image_gen_client


def get_image_edit_client() -> openai.AsyncOpenAI:
    """Return the shared async image-edit client, building it once."""
    global _image_edit_client
    if _image_edit_client is None:
        from vibe_bot.config import IMAGE_EDIT_ENDPOINT, IMAGE_EDIT_ENDPOINT_KEY

        _image_edit_client = openai.AsyncOpenAI(
            base_url=IMAGE_EDIT_ENDPOINT, api_key=IMAGE_EDIT_ENDPOINT_KEY
        )
    return _image_edit_client


async def image_generation(
    prompt: str,
    *,
    model: str = "gen",
    n: int = 1,
    size: str = "1024x1024",
) -> str:
    """Generate an image; return base64 data ("" on failure)."""
    client = get_image_gen_client()
    try:
        response = await client.images.generate(
            prompt=prompt,
            n=n,
            size=size,
            model=model,
            timeout=300.0,
        )
    except openai.OpenAIError as e:
        logger.warning("Image generation failed: %s", e)
        return ""
    if response.data:
        return response.data[0].b64_json or ""
    return ""


async def image_edit(
    image: BufferedReader | BytesIO | list[BufferedReader] | list[BytesIO],
    prompt: str,
    *,
    model: str = "edit",
    n: int = 1,
) -> str:
    """Edit an image; return base64 data ("" on failure)."""
    client = get_image_edit_client()
    try:
        response = await client.images.edit(
            image=image,
            prompt=prompt,
            n=n,
            size="768x768",
            model=model,
        )
    except openai.OpenAIError as e:
        logger.warning("Image edit failed: %s", e)
        return ""
    if response.data:
        return response.data[0].b64_json or ""
    return ""