from __future__ import annotations

from typing import TYPE_CHECKING, List, Optional, Sequence

from .._resource import AsyncAPIResource, SyncAPIResource
from .._types import NOT_GIVEN, FileInput, NotGiven, is_given
from .._utils import strip_not_given
from ..types.create_response import CreateResponse
from ..types.generation import Generation

if TYPE_CHECKING:
    from .files import AsyncFiles, Files


class Images(SyncAPIResource):
    @property
    def _files(self) -> Files:
        return self._client.files  # type: ignore[return-value]

    def submit(
        self,
        *,
        model: str,
        prompt: str,
        size: str | NotGiven = NOT_GIVEN,
        n: int | NotGiven = NOT_GIVEN,
        quality: str | NotGiven = NOT_GIVEN,
        output_format: str | NotGiven = NOT_GIVEN,
        resolution: str | NotGiven = NOT_GIVEN,
        aspect_ratio: str | NotGiven = NOT_GIVEN,
        reference_images: List[FileInput] | NotGiven = NOT_GIVEN,
        mask: FileInput | NotGiven = NOT_GIVEN,
        context: str | NotGiven = NOT_GIVEN,
        **extra_body: object,
    ) -> CreateResponse:
        body, _ = self._build_body(
            model=model, prompt=prompt, size=size, n=n,
            quality=quality, output_format=output_format,
            resolution=resolution, aspect_ratio=aspect_ratio,
            reference_images=reference_images, mask=mask, context=context,
            **extra_body,
        )
        return self._post("/generateimage", cast_to=CreateResponse, body=body)

    def generate(
        self,
        *,
        model: str,
        prompt: str,
        size: str | NotGiven = NOT_GIVEN,
        n: int | NotGiven = NOT_GIVEN,
        quality: str | NotGiven = NOT_GIVEN,
        output_format: str | NotGiven = NOT_GIVEN,
        resolution: str | NotGiven = NOT_GIVEN,
        aspect_ratio: str | NotGiven = NOT_GIVEN,
        reference_images: List[FileInput] | NotGiven = NOT_GIVEN,
        mask: FileInput | NotGiven = NOT_GIVEN,
        context: str | NotGiven = NOT_GIVEN,
        poll_interval: float = 2.0,
        timeout: float = 120.0,
        **extra_body: object,
    ) -> Generation:
        body, temp_ids = self._build_body(
            model=model, prompt=prompt, size=size, n=n,
            quality=quality, output_format=output_format,
            resolution=resolution, aspect_ratio=aspect_ratio,
            reference_images=reference_images, mask=mask, context=context,
            **extra_body,
        )
        try:
            resp = self._post("/generateimage", cast_to=CreateResponse, body=body)
            return self._client.generations.wait(  # type: ignore[union-attr]
                resp.request_id,
                poll_interval=poll_interval,
                timeout=timeout,
            )
        finally:
            if temp_ids:
                self._files.cleanup(temp_ids)

    def _build_body(self, **kwargs: object) -> tuple[dict, list[str]]:
        temp: list[str] = []
        reference_images = kwargs.pop("reference_images")
        mask = kwargs.pop("mask")
        context = kwargs.pop("context")
        body: dict = strip_not_given({**kwargs, "context": context})  # type: ignore[arg-type]
        if is_given(reference_images):
            body["reference_images"] = self._files.resolve_many(reference_images, _temp=temp)  # type: ignore[arg-type]
        if is_given(mask):
            body["mask"] = self._files.resolve(mask, _temp=temp)  # type: ignore[arg-type]
        return body, temp


class AsyncImages(AsyncAPIResource):
    @property
    def _files(self) -> AsyncFiles:
        return self._client.files  # type: ignore[return-value]

    async def submit(
        self,
        *,
        model: str,
        prompt: str,
        size: str | NotGiven = NOT_GIVEN,
        n: int | NotGiven = NOT_GIVEN,
        quality: str | NotGiven = NOT_GIVEN,
        output_format: str | NotGiven = NOT_GIVEN,
        resolution: str | NotGiven = NOT_GIVEN,
        aspect_ratio: str | NotGiven = NOT_GIVEN,
        reference_images: List[FileInput] | NotGiven = NOT_GIVEN,
        mask: FileInput | NotGiven = NOT_GIVEN,
        context: str | NotGiven = NOT_GIVEN,
        **extra_body: object,
    ) -> CreateResponse:
        body, _ = await self._build_body(
            model=model, prompt=prompt, size=size, n=n,
            quality=quality, output_format=output_format,
            resolution=resolution, aspect_ratio=aspect_ratio,
            reference_images=reference_images, mask=mask, context=context,
            **extra_body,
        )
        return await self._post("/generateimage", cast_to=CreateResponse, body=body)

    async def generate(
        self,
        *,
        model: str,
        prompt: str,
        size: str | NotGiven = NOT_GIVEN,
        n: int | NotGiven = NOT_GIVEN,
        quality: str | NotGiven = NOT_GIVEN,
        output_format: str | NotGiven = NOT_GIVEN,
        resolution: str | NotGiven = NOT_GIVEN,
        aspect_ratio: str | NotGiven = NOT_GIVEN,
        reference_images: List[FileInput] | NotGiven = NOT_GIVEN,
        mask: FileInput | NotGiven = NOT_GIVEN,
        context: str | NotGiven = NOT_GIVEN,
        poll_interval: float = 2.0,
        timeout: float = 120.0,
        **extra_body: object,
    ) -> Generation:
        body, temp_ids = await self._build_body(
            model=model, prompt=prompt, size=size, n=n,
            quality=quality, output_format=output_format,
            resolution=resolution, aspect_ratio=aspect_ratio,
            reference_images=reference_images, mask=mask, context=context,
            **extra_body,
        )
        try:
            resp = await self._post("/generateimage", cast_to=CreateResponse, body=body)
            return await self._client.generations.wait(  # type: ignore[union-attr]
                resp.request_id,
                poll_interval=poll_interval,
                timeout=timeout,
            )
        finally:
            if temp_ids:
                await self._files.cleanup(temp_ids)

    async def _build_body(self, **kwargs: object) -> tuple[dict, list[str]]:
        temp: list[str] = []
        reference_images = kwargs.pop("reference_images")
        mask = kwargs.pop("mask")
        context = kwargs.pop("context")
        body: dict = strip_not_given({**kwargs, "context": context})  # type: ignore[arg-type]
        if is_given(reference_images):
            body["reference_images"] = await self._files.resolve_many(reference_images, _temp=temp)  # type: ignore[arg-type]
        if is_given(mask):
            body["mask"] = await self._files.resolve(mask, _temp=temp)  # type: ignore[arg-type]
        return body, temp
