"""Image generation tool using generated API client.

Provides both synchronous and asynchronous image generation
with support for model aliases, quality/style options,
async job polling, and preset management.

Example:
    from sdkrouter import SDKRouter

    client = SDKRouter(api_key="...")

    # Sync generation
    result = client.image_gen.generate(
        prompt="A sunset over mountains",
        model="@balanced",
        size="1024x1024",
        crop_to="1200x630",
    )
    print(f"Image URL: {result.image_cdn_url}")

    # Async generation with polling
    job = client.image_gen.generate_async(prompt="...")
    result = client.image_gen.wait_for_completion(job.id)

    # Preset management
    preset = client.image_gen.create_preset(
        name="OG Image",
        slug="og-image",
        size="1792x1024",
        crop_to="1200x630",
    )
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import TYPE_CHECKING

from .._api.client import (
    BaseResource,
    AsyncBaseResource,
)
from .._api.generated.image_gen.image_gen__api__image_gen.sync_client import (
    SyncImageGenImageGenAPI,
)
from .._api.generated.image_gen.image_gen__api__image_gen.client import (
    ImageGenImageGenAPI,
)
from .._api.generated.image_gen.image_gen__api__image_gen.models import (
    ImageGenerateRequestRequest,
    ImageGenerateResponse,
    AsyncGenerateResponse,
    JobStatus,
    ImageGenerationList,
    ImageGenerationDetail,
    ImageGenOptions,
    ChoiceItem,
    ImageGenPreset,
    ImageGenPresetCreateRequest,
    PatchedImageGenPresetCreateRequest,
)
from .._api.generated.image_gen.enums import (
    ImageGenerateResponseStatus,
    ImageGenPresetQuality,
    ImageGenPresetStyle,
)

# Quality/style enums are shared with presets
ImageGenerateRequestRequestQuality = ImageGenPresetQuality
ImageGenerateRequestRequestStyle = ImageGenPresetStyle
from ..exceptions import api_error_handler, async_api_error_handler, APIError

if TYPE_CHECKING:
    from .._config import SDKConfig

logger = logging.getLogger(__name__)


class ImageGenResource(BaseResource):
    """Image generation tool (sync).

    Uses generated SyncImageGenImageGenAPI client.

    Example:
        result = client.image_gen.generate(
            prompt="A beautiful sunset",
            model="@balanced",
            quality="hd",
            crop_to="1200x630",
        )
        print(f"Generated: {result.image_cdn_url}")
    """

    def __init__(self, config: "SDKConfig"):
        super().__init__(config)
        self._api = SyncImageGenImageGenAPI(self._http_client)

    @api_error_handler
    def generate(
        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        model: str | None = None,
        size: str | None = None,
        quality: ImageGenerateRequestRequestQuality | str | None = None,
        style: ImageGenerateRequestRequestStyle | str | None = None,
        crop_to: str | None = None,
        preset: str | None = None,
    ) -> ImageGenerateResponse:
        """Generate an image synchronously.

        Args:
            prompt: Text description of the image to generate (max 4000 chars)
            negative_prompt: What to avoid in the image (max 1000 chars)
            model: Model to use. Supports aliases: @cheap, @balanced, @smart, @fast
            size: Image size (e.g., "1024x1024", "1792x1024")
            quality: Quality level ("standard" or "hd")
            style: Style ("natural" or "vivid")
            crop_to: Crop output to this size (e.g., "1200x630")
            preset: Preset slug to apply defaults from

        Returns:
            ImageGenerateResponse with image URLs and metadata

        Raises:
            ValidationError: If request parameters are invalid
            RateLimitError: If rate limit is exceeded
            APIError: For other API errors
        """
        logger.debug("Image generate: prompt=%s, model=%s", prompt[:50], model)

        kwargs: dict = {"prompt": prompt}
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if preset is not None:
            kwargs["preset"] = preset

        request = ImageGenerateRequestRequest(**kwargs)
        result = self._api.generate_create(request)

        logger.info(
            "Image generated: id=%s, cost=$%s",
            result.id,
            result.cost_usd or "0",
        )
        return result

    @api_error_handler
    def generate_async(
        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        model: str | None = None,
        size: str | None = None,
        quality: ImageGenerateRequestRequestQuality | str | None = None,
        style: ImageGenerateRequestRequestStyle | str | None = None,
        crop_to: str | None = None,
        preset: str | None = None,
    ) -> AsyncGenerateResponse:
        """Start async image generation.

        Returns immediately with a job_id for polling.

        Args:
            prompt: Text description of the image to generate
            negative_prompt: What to avoid in the image
            model: Model to use (supports aliases)
            size: Image size
            quality: Quality level
            style: Style
            crop_to: Crop output to this size (e.g., "1200x630")
            preset: Preset slug to apply defaults from

        Returns:
            AsyncGenerateResponse with id and job_id for polling

        Example:
            job = client.image_gen.generate_async(prompt="...")
            result = client.image_gen.wait_for_completion(job.id)
        """
        logger.debug("Image generate async: prompt=%s", prompt[:50])

        kwargs: dict = {"prompt": prompt}
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if preset is not None:
            kwargs["preset"] = preset

        request = ImageGenerateRequestRequest(**kwargs)
        result = self._api.generate_async_create(request)

        logger.info("Image generation queued: id=%s, job_id=%s", result.id, result.job_id)
        return result

    @api_error_handler
    def get_status(self, generation_id: int | str) -> JobStatus:
        """Get status of an async generation job."""
        return self._api.status_retrieve(str(generation_id))

    @api_error_handler
    def get(self, generation_id: int | str) -> ImageGenerationDetail:
        """Get details of a generation."""
        return self._api.retrieve(str(generation_id))

    @api_error_handler
    def list(
        self,
        *,
        ordering: str | None = None,
        search: str | None = None,
    ) -> list[ImageGenerationList]:
        """List generation history."""
        return self._api.list(ordering=ordering, search=search)

    @api_error_handler
    def options(self) -> ImageGenOptions:
        """Get available generation options."""
        return self._api.options_retrieve()

    def wait_for_completion(
        self,
        generation_id: int | str,
        *,
        timeout: float = 300.0,
        poll_interval: float = 2.0,
    ) -> ImageGenerationDetail:
        """Wait for async generation to complete.

        Raises:
            TimeoutError: If generation doesn't complete within timeout
            APIError: If generation fails
        """
        start_time = time.monotonic()

        while True:
            elapsed = time.monotonic() - start_time
            if elapsed > timeout:
                raise TimeoutError(
                    f"Image generation {generation_id} did not complete within {timeout}s"
                )

            status = self.get_status(generation_id)
            logger.debug(
                "Generation %s status: %s (elapsed: %.1fs)",
                generation_id,
                status.status,
                elapsed,
            )

            if status.status == ImageGenerateResponseStatus.COMPLETED:
                return self.get(generation_id)
            elif status.status in (
                ImageGenerateResponseStatus.FAILED,
                ImageGenerateResponseStatus.CANCELLED,
            ):
                error_msg = status.error or f"Generation {status.status}"
                raise APIError(error_msg)

            time.sleep(poll_interval)

    # --- Preset CRUD ---

    @api_error_handler
    def list_presets(
        self,
        *,
        ordering: str | None = None,
        search: str | None = None,
    ) -> list[ImageGenPreset]:
        """List image generation presets for the current API key."""
        return self._api.presets_list(ordering=ordering, search=search)

    @api_error_handler
    def get_preset(self, preset_id: str) -> ImageGenPreset:
        """Get a preset by ID."""
        return self._api.presets_retrieve(preset_id)

    @api_error_handler
    def create_preset(
        self,
        *,
        name: str,
        slug: str,
        model: str | None = None,
        size: str | None = None,
        crop_to: str | None = None,
        quality: ImageGenPresetQuality | str | None = None,
        style: ImageGenPresetStyle | str | None = None,
        negative_prompt: str | None = None,
        metadata: dict | None = None,
    ) -> ImageGenPreset:
        """Create an image generation preset.

        Args:
            name: Human-readable preset name
            slug: URL-safe identifier (unique per API key)
            model: Model ID or alias (empty = use request default)
            size: Image size (empty = use request default)
            crop_to: Crop target size (empty = no crop)
            quality: Image quality (empty = use request default)
            style: Image style (empty = use request default)
            negative_prompt: Default negative prompt
            metadata: Arbitrary extra metadata
        """
        kwargs: dict = {"name": name, "slug": slug}
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if metadata is not None:
            kwargs["metadata"] = metadata

        data = ImageGenPresetCreateRequest(**kwargs)
        return self._api.presets_create(data)

    @api_error_handler
    def update_preset(
        self,
        preset_id: str,
        *,
        name: str,
        slug: str,
        model: str | None = None,
        size: str | None = None,
        crop_to: str | None = None,
        quality: ImageGenPresetQuality | str | None = None,
        style: ImageGenPresetStyle | str | None = None,
        negative_prompt: str | None = None,
        metadata: dict | None = None,
    ) -> ImageGenPreset:
        """Replace a preset (full update)."""
        kwargs: dict = {"name": name, "slug": slug}
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if metadata is not None:
            kwargs["metadata"] = metadata

        data = ImageGenPresetCreateRequest(**kwargs)
        return self._api.presets_update(preset_id, data)

    @api_error_handler
    def patch_preset(
        self,
        preset_id: str,
        *,
        name: str | None = None,
        slug: str | None = None,
        model: str | None = None,
        size: str | None = None,
        crop_to: str | None = None,
        quality: ImageGenPresetQuality | str | None = None,
        style: ImageGenPresetStyle | str | None = None,
        negative_prompt: str | None = None,
        metadata: dict | None = None,
    ) -> ImageGenPreset:
        """Partially update a preset (PATCH)."""
        kwargs: dict = {}
        if name is not None:
            kwargs["name"] = name
        if slug is not None:
            kwargs["slug"] = slug
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if metadata is not None:
            kwargs["metadata"] = metadata

        data = PatchedImageGenPresetCreateRequest(**kwargs)
        return self._api.presets_partial_update(preset_id, data)

    @api_error_handler
    def delete_preset(self, preset_id: str) -> None:
        """Delete a preset."""
        return self._api.presets_destroy(preset_id)


class AsyncImageGenResource(AsyncBaseResource):
    """Image generation tool (async).

    Uses generated ImageGenImageGenAPI client.

    Example:
        result = await client.image_gen.generate(
            prompt="A beautiful sunset",
            model="@balanced",
            crop_to="1200x630",
        )
        print(f"Generated: {result.image_cdn_url}")
    """

    def __init__(self, config: "SDKConfig"):
        super().__init__(config)
        self._api = ImageGenImageGenAPI(self._http_client)

    @async_api_error_handler
    async def generate(
        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        model: str | None = None,
        size: str | None = None,
        quality: ImageGenerateRequestRequestQuality | str | None = None,
        style: ImageGenerateRequestRequestStyle | str | None = None,
        crop_to: str | None = None,
        preset: str | None = None,
    ) -> ImageGenerateResponse:
        """Generate an image asynchronously.

        Args:
            prompt: Text description of the image to generate (max 4000 chars)
            negative_prompt: What to avoid in the image (max 1000 chars)
            model: Model to use. Supports aliases: @cheap, @balanced, @smart, @fast
            size: Image size (e.g., "1024x1024", "1792x1024")
            quality: Quality level ("standard" or "hd")
            style: Style ("natural" or "vivid")
            crop_to: Crop output to this size (e.g., "1200x630")
            preset: Preset slug to apply defaults from
        """
        logger.debug("Image generate async: prompt=%s", prompt[:50])

        kwargs: dict = {"prompt": prompt}
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if preset is not None:
            kwargs["preset"] = preset

        request = ImageGenerateRequestRequest(**kwargs)
        result = await self._api.generate_create(request)

        logger.info("Image generated: id=%s, cost=$%s", result.id, result.cost_usd or "0")
        return result

    @async_api_error_handler
    async def generate_async(
        self,
        prompt: str,
        *,
        negative_prompt: str | None = None,
        model: str | None = None,
        size: str | None = None,
        quality: ImageGenerateRequestRequestQuality | str | None = None,
        style: ImageGenerateRequestRequestStyle | str | None = None,
        crop_to: str | None = None,
        preset: str | None = None,
    ) -> AsyncGenerateResponse:
        """Start async image generation (returns immediately).

        Args:
            prompt: Text description of the image to generate
            crop_to: Crop output to this size (e.g., "1200x630")
            preset: Preset slug to apply defaults from
        """
        kwargs: dict = {"prompt": prompt}
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if preset is not None:
            kwargs["preset"] = preset

        request = ImageGenerateRequestRequest(**kwargs)
        result = await self._api.generate_async_create(request)

        logger.info("Image generation queued: id=%s, job_id=%s", result.id, result.job_id)
        return result

    @async_api_error_handler
    async def get_status(self, generation_id: int | str) -> JobStatus:
        """Get status of an async generation job."""
        return await self._api.status_retrieve(str(generation_id))

    @async_api_error_handler
    async def get(self, generation_id: int | str) -> ImageGenerationDetail:
        """Get details of a generation."""
        return await self._api.retrieve(str(generation_id))

    @async_api_error_handler
    async def list(
        self,
        *,
        ordering: str | None = None,
        search: str | None = None,
    ) -> list[ImageGenerationList]:
        """List generation history."""
        return await self._api.list(ordering=ordering, search=search)

    @async_api_error_handler
    async def options(self) -> ImageGenOptions:
        """Get available generation options."""
        return await self._api.options_retrieve()

    async def wait_for_completion(
        self,
        generation_id: int | str,
        *,
        timeout: float = 300.0,
        poll_interval: float = 2.0,
    ) -> ImageGenerationDetail:
        """Wait for async generation to complete.

        Raises:
            TimeoutError: If generation doesn't complete within timeout
            APIError: If generation fails
        """
        start_time = time.monotonic()

        while True:
            elapsed = time.monotonic() - start_time
            if elapsed > timeout:
                raise TimeoutError(
                    f"Image generation {generation_id} did not complete within {timeout}s"
                )

            status = await self.get_status(generation_id)
            logger.debug(
                "Generation %s status: %s (elapsed: %.1fs)",
                generation_id,
                status.status,
                elapsed,
            )

            if status.status == ImageGenerateResponseStatus.COMPLETED:
                return await self.get(generation_id)
            elif status.status in (
                ImageGenerateResponseStatus.FAILED,
                ImageGenerateResponseStatus.CANCELLED,
            ):
                error_msg = status.error or f"Generation {status.status}"
                raise APIError(error_msg)

            await asyncio.sleep(poll_interval)

    # --- Preset CRUD ---

    @async_api_error_handler
    async def list_presets(
        self,
        *,
        ordering: str | None = None,
        search: str | None = None,
    ) -> list[ImageGenPreset]:
        """List image generation presets for the current API key."""
        return await self._api.presets_list(ordering=ordering, search=search)

    @async_api_error_handler
    async def get_preset(self, preset_id: str) -> ImageGenPreset:
        """Get a preset by ID."""
        return await self._api.presets_retrieve(preset_id)

    @async_api_error_handler
    async def create_preset(
        self,
        *,
        name: str,
        slug: str,
        model: str | None = None,
        size: str | None = None,
        crop_to: str | None = None,
        quality: ImageGenPresetQuality | str | None = None,
        style: ImageGenPresetStyle | str | None = None,
        negative_prompt: str | None = None,
        metadata: dict | None = None,
    ) -> ImageGenPreset:
        """Create an image generation preset."""
        kwargs: dict = {"name": name, "slug": slug}
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if metadata is not None:
            kwargs["metadata"] = metadata

        data = ImageGenPresetCreateRequest(**kwargs)
        return await self._api.presets_create(data)

    @async_api_error_handler
    async def update_preset(
        self,
        preset_id: str,
        *,
        name: str,
        slug: str,
        model: str | None = None,
        size: str | None = None,
        crop_to: str | None = None,
        quality: ImageGenPresetQuality | str | None = None,
        style: ImageGenPresetStyle | str | None = None,
        negative_prompt: str | None = None,
        metadata: dict | None = None,
    ) -> ImageGenPreset:
        """Replace a preset (full update)."""
        kwargs: dict = {"name": name, "slug": slug}
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if metadata is not None:
            kwargs["metadata"] = metadata

        data = ImageGenPresetCreateRequest(**kwargs)
        return await self._api.presets_update(preset_id, data)

    @async_api_error_handler
    async def patch_preset(
        self,
        preset_id: str,
        *,
        name: str | None = None,
        slug: str | None = None,
        model: str | None = None,
        size: str | None = None,
        crop_to: str | None = None,
        quality: ImageGenPresetQuality | str | None = None,
        style: ImageGenPresetStyle | str | None = None,
        negative_prompt: str | None = None,
        metadata: dict | None = None,
    ) -> ImageGenPreset:
        """Partially update a preset (PATCH)."""
        kwargs: dict = {}
        if name is not None:
            kwargs["name"] = name
        if slug is not None:
            kwargs["slug"] = slug
        if model is not None:
            kwargs["model"] = model
        if size is not None:
            kwargs["size"] = size
        if crop_to is not None:
            kwargs["crop_to"] = crop_to
        if quality is not None:
            kwargs["quality"] = quality
        if style is not None:
            kwargs["style"] = style
        if negative_prompt is not None:
            kwargs["negative_prompt"] = negative_prompt
        if metadata is not None:
            kwargs["metadata"] = metadata

        data = PatchedImageGenPresetCreateRequest(**kwargs)
        return await self._api.presets_partial_update(preset_id, data)

    @async_api_error_handler
    async def delete_preset(self, preset_id: str) -> None:
        """Delete a preset."""
        return await self._api.presets_destroy(preset_id)


__all__ = [
    # Resources
    "ImageGenResource",
    "AsyncImageGenResource",
    # Request/Response models
    "ImageGenerateRequestRequest",
    "ImageGenerateResponse",
    "AsyncGenerateResponse",
    "JobStatus",
    "ImageGenerationList",
    "ImageGenerationDetail",
    "ImageGenOptions",
    "ChoiceItem",
    # Preset models
    "ImageGenPreset",
    "ImageGenPresetCreateRequest",
    "PatchedImageGenPresetCreateRequest",
    # Enums
    "ImageGenerateRequestRequestQuality",
    "ImageGenerateRequestRequestStyle",
    "ImageGenerateResponseStatus",
    "ImageGenPresetQuality",
    "ImageGenPresetStyle",
]
