"""Nano Banana (Gemini) image generation resource."""
from __future__ import annotations

import asyncio
import logging
from typing import TYPE_CHECKING

from ..types.images import (
    ImageModel,
    ImageQuality,
    ImageAspectRatio,
    ImageStatus,
    ImageGenerateParams,
    ImageGenerateResponse,
    ImageStatusResponse,
)
from .._errors import GenerationTimeoutError

if TYPE_CHECKING:
    from ..client import RunBlob

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

# Default polling intervals
_POLL_INITIAL_DELAY = 1.0    # First check after 1s
_POLL_MAX_DELAY = 5.0        # Cap at 5s between checks
_POLL_BACKOFF_FACTOR = 1.5   # Multiply delay each iteration
_POLL_DEFAULT_TIMEOUT = 300.0  # 5 minutes max wait


class NanoBanana:
    """Nano Banana (Gemini) image generation.

    Available via ``client.nanobanana``.
    """

    _client: RunBlob

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

    async def generate(
        self,
        *,
        prompt: str,
        model: ImageModel | str = ImageModel.STANDARD,
        quality: ImageQuality | str = ImageQuality.STANDARD,
        aspect_ratio: ImageAspectRatio | str = ImageAspectRatio.AUTO,
        temperature: float = 0.7,
        images: list[str] | None = None,
        callback_url: str | None = None,
    ) -> ImageGenerateResponse:
        """Submit an image generation task.

        Returns immediately with a task_uuid for tracking.

        Args:
            prompt: Text description (1-4000 chars).
            model: "standard", "pro", or "v2".
            quality: "standard", "2k", or "4k" (pro/v2 only).
            aspect_ratio: Image aspect ratio (e.g., "16:9", "1:1").
            temperature: Creativity 0.0-1.0 (default 0.7).
            images: Input images for editing (base64 or URLs).
            callback_url: Webhook URL for completion notification.

        Returns:
            ImageGenerateResponse with task_uuid, status, and price.

        Raises:
            AuthenticationError: Invalid API key.
            InsufficientCreditsError: Not enough balance.
            APIStatusError: Invalid parameters.
        """
        # Coerce string enums
        model_enum = ImageModel(model) if isinstance(model, str) else model
        quality_enum = ImageQuality(quality) if isinstance(quality, str) else quality
        ar_enum = ImageAspectRatio(aspect_ratio) if isinstance(aspect_ratio, str) else aspect_ratio

        params = ImageGenerateParams(
            prompt=prompt,
            model=model_enum,
            quality=quality_enum,
            aspect_ratio=ar_enum,
            temperature=temperature,
            images=images,
            callback_url=callback_url,
        )

        return await self._client.request(
            "POST",
            "/v1/gemini/generate",
            response_model=ImageGenerateResponse,
            body=params.to_api_body(),
        )

    async def get_status(
        self,
        task_uuid: str,
        *,
        use_cache: bool = False,
    ) -> ImageStatusResponse:
        """Get the current status of a generation task.

        Args:
            task_uuid: The UUID returned from generate().
            use_cache: Use cached response if available (default False).

        Returns:
            ImageStatusResponse with status, prompt, result_image_url, message.
        """
        return await self._client.request(
            "GET",
            f"/v1/gemini/generations/{task_uuid}",
            response_model=ImageStatusResponse,
            use_cache=use_cache,
        )

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

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

        Returns:
            List of ImageStatusResponse objects.
        """
        response = await self._client.request_raw(
            "GET",
            "/v1/gemini/generations",
            params={"limit": limit, "offset": offset},
        )
        items = response.json()
        return [ImageStatusResponse.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,
    ) -> ImageStatusResponse:
        """Poll until generation reaches a terminal state (completed/failed).

        Uses exponential backoff to reduce API calls during long generations.

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

        Returns:
            ImageStatusResponse with final status.

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

        logger.info("Polling generation %s (timeout=%.0fs)", task_uuid, timeout)

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

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

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

            # Sleep with backoff, but don't exceed remaining time
            actual_delay = min(delay, remaining)
            logger.debug(
                "Generation %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: ImageModel | str = ImageModel.STANDARD,
        quality: ImageQuality | str = ImageQuality.STANDARD,
        aspect_ratio: ImageAspectRatio | str = ImageAspectRatio.AUTO,
        temperature: float = 0.7,
        images: list[str] | None = None,
        callback_url: str | None = None,
        timeout: float = _POLL_DEFAULT_TIMEOUT,
    ) -> ImageStatusResponse:
        """Submit a generation and wait for the result.

        Convenience method combining generate() + wait_for_result().

        Returns:
            ImageStatusResponse with final status and result_image_url.
        """
        task = await self.generate(
            prompt=prompt,
            model=model,
            quality=quality,
            aspect_ratio=aspect_ratio,
            temperature=temperature,
            images=images,
            callback_url=callback_url,
        )
        logger.info(
            "Created generation %s (price=%s), waiting for result...",
            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: ImageModel | str = ImageModel.STANDARD,
        quality: ImageQuality | str = ImageQuality.STANDARD,
        aspect_ratio: ImageAspectRatio | str = ImageAspectRatio.AUTO,
        temperature: float = 0.7,
        concurrency: int = 10,
    ) -> list[ImageGenerateResponse]:
        """Submit multiple generations concurrently.

        Uses a semaphore to control concurrency and avoid overwhelming the API.

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

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

        async def _generate_one(prompt: str) -> ImageGenerateResponse:
            async with sem:
                return await self.generate(
                    prompt=prompt,
                    model=model,
                    quality=quality,
                    aspect_ratio=aspect_ratio,
                    temperature=temperature,
                )

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

    async def generate_batch_and_wait(
        self,
        prompts: list[str],
        *,
        model: ImageModel | str = ImageModel.STANDARD,
        quality: ImageQuality | str = ImageQuality.STANDARD,
        aspect_ratio: ImageAspectRatio | str = ImageAspectRatio.AUTO,
        temperature: float = 0.7,
        concurrency: int = 10,
        timeout: float = _POLL_DEFAULT_TIMEOUT,
    ) -> list[ImageStatusResponse]:
        """Submit multiple generations and wait for all results.

        Args:
            prompts: List of text prompts.
            concurrency: Max concurrent submissions (default 10).
            timeout: Per-generation poll timeout (default 300s).

        Returns:
            List of ImageStatusResponse in same order as prompts.
        """
        tasks = await self.generate_batch(
            prompts,
            model=model,
            quality=quality,
            aspect_ratio=aspect_ratio,
            temperature=temperature,
            concurrency=concurrency,
        )
        return list(
            await asyncio.gather(
                *[self.wait_for_result(t.task_uuid, timeout=timeout) for t in tasks]
            )
        )
