"""ChatGPT image generation resource."""
from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING

from ..types.chatgpt import (
    ChatGPTModel,
    ChatGPTAspectRatio,
    ChatGPTStatus,
    ChatGPTGenerateParams,
    ChatGPTGenerateResponse,
    ChatGPTStatusResponse,
)
from .._errors import GenerationTimeoutError

if TYPE_CHECKING:
    from ..client import RunBlob

logger = logging.getLogger("runblob.chatgpt")

_POLL_INITIAL_DELAY = 2.0
_POLL_MAX_DELAY = 10.0
_POLL_BACKOFF_FACTOR = 1.5
_POLL_DEFAULT_TIMEOUT = 300.0


class ChatGPT:
    """ChatGPT image generation.

    Available via ``client.chatgpt``.
    """

    _client: RunBlob

    def __init__(self, client: RunBlob) -> None:
        self._client = client

    async def generate(
        self,
        *,
        prompt: str,
        model: ChatGPTModel | str = ChatGPTModel.GPT_5_2,
        aspect_ratio: ChatGPTAspectRatio | str | None = None,
        images: list[str] | None = None,
        callback_url: str | None = None,
    ) -> ChatGPTGenerateResponse:
        """Generate an image with ChatGPT.

        Supports text-to-image and image editing (up to 4 input images).

        Args:
            prompt: Text description (1-4000 chars).
            model: ChatGPT model (default "gpt-5-2").
            aspect_ratio: "1:1", "4:3", "3:4", "3:2", or "2:3".
            images: Input images for editing — URLs, base64, or data URIs (max 4, 20MB each).
            callback_url: Webhook URL for completion notification.

        Returns:
            ChatGPTGenerateResponse with task_uuid, status, price.
        """
        model_enum = ChatGPTModel(model) if isinstance(model, str) else model
        ar_enum = ChatGPTAspectRatio(aspect_ratio) if isinstance(aspect_ratio, str) and aspect_ratio else aspect_ratio

        params = ChatGPTGenerateParams(
            prompt=prompt,
            model=model_enum,
            aspect_ratio=ar_enum if isinstance(ar_enum, ChatGPTAspectRatio) else None,
            images=images,
            callback_url=callback_url,
        )

        return await self._client.request(
            "POST",
            "/v1/chatgpt-images/generate",
            response_model=ChatGPTGenerateResponse,
            body=params.to_api_body(),
        )

    async def get_status(self, task_uuid: str) -> ChatGPTStatusResponse:
        """Get image generation status.

        Args:
            task_uuid: The UUID returned from generate().

        Returns:
            ChatGPTStatusResponse with status, result_image_url on completion.
        """
        return await self._client.request(
            "GET",
            f"/v1/chatgpt-images/generations/{task_uuid}",
            response_model=ChatGPTStatusResponse,
        )

    async def list(
        self,
        *,
        limit: int = 50,
        offset: int = 0,
    ) -> list[ChatGPTStatusResponse]:
        """List recent generations.

        Args:
            limit: Max results, 1-100 (default 50).
            offset: Skip N results for pagination (default 0).

        Returns:
            List of ChatGPTStatusResponse objects.
        """
        response = await self._client.request_raw(
            "GET",
            "/v1/chatgpt-images/generations",
            params={"limit": limit, "offset": offset},
        )
        items = response.json()
        return [ChatGPTStatusResponse.model_validate(item) for item in items]

    async def wait_for_result(
        self,
        task_uuid: str,
        *,
        timeout: float = _POLL_DEFAULT_TIMEOUT,
        poll_interval: float = _POLL_INITIAL_DELAY,
        max_poll_interval: float = _POLL_MAX_DELAY,
        backoff_factor: float = _POLL_BACKOFF_FACTOR,
    ) -> ChatGPTStatusResponse:
        """Poll until image generation completes or fails.

        Args:
            task_uuid: The UUID to poll.
            timeout: Max seconds to wait (default 300).
            poll_interval: Initial delay between polls (default 2s).
            max_poll_interval: Max delay between polls (default 10s).
            backoff_factor: Multiply interval each iteration (default 1.5).

        Returns:
            ChatGPTStatusResponse with final status.

        Raises:
            GenerationTimeoutError: If timeout exceeded.
        """
        deadline = asyncio.get_event_loop().time() + timeout
        delay = poll_interval

        while True:
            status = await self.get_status(task_uuid)

            if status.is_terminal:
                logger.info("ChatGPT %s finished: %s", task_uuid, status.status.value)
                return status

            remaining = deadline - asyncio.get_event_loop().time()
            if remaining <= 0:
                raise GenerationTimeoutError(task_uuid, timeout)

            actual_delay = min(delay, remaining)
            logger.debug("ChatGPT %s status=%s, polling in %.1fs", task_uuid, status.status.value, actual_delay)
            await asyncio.sleep(actual_delay)
            delay = min(delay * backoff_factor, max_poll_interval)

    async def generate_and_wait(
        self,
        *,
        prompt: str,
        model: ChatGPTModel | str = ChatGPTModel.GPT_5_2,
        aspect_ratio: ChatGPTAspectRatio | str | None = None,
        images: list[str] | None = None,
        callback_url: str | None = None,
        timeout: float = _POLL_DEFAULT_TIMEOUT,
    ) -> ChatGPTStatusResponse:
        """Generate image and wait for result.

        Returns:
            ChatGPTStatusResponse with final status and result_image_url.
        """
        task = await self.generate(
            prompt=prompt, model=model, aspect_ratio=aspect_ratio,
            images=images, callback_url=callback_url,
        )
        logger.info("Created ChatGPT %s (price=%s), waiting...", task.task_uuid, task.price)
        return await self.wait_for_result(task.task_uuid, timeout=timeout)

    async def generate_batch(
        self,
        prompts: list[str],
        *,
        model: ChatGPTModel | str = ChatGPTModel.GPT_5_2,
        aspect_ratio: ChatGPTAspectRatio | str | None = None,
        concurrency: int = 10,
    ) -> list[ChatGPTGenerateResponse]:
        """Submit multiple image generations concurrently.

        Args:
            prompts: List of text prompts.
            model: Model for all generations.
            aspect_ratio: Aspect ratio for all.
            concurrency: Max concurrent requests (default 10).

        Returns:
            List of ChatGPTGenerateResponse in same order as prompts.
        """
        sem = asyncio.Semaphore(concurrency)

        async def _gen(prompt: str) -> ChatGPTGenerateResponse:
            async with sem:
                return await self.generate(
                    prompt=prompt, model=model, aspect_ratio=aspect_ratio,
                )

        return list(await asyncio.gather(*[_gen(p) for p in prompts]))
