from __future__ import annotations

from typing import TYPE_CHECKING, Any

from ..._endpoints import Prompts
from ..._internal._schemas import APIEnvelope
from ..._internal._schemas.prompts import (
    PromptRecommendation,
)
from .._base import BaseResource

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from ..._async_client import AsyncClient
else:
    AsyncClient = Any


class AsyncPromptsResource(BaseResource[AsyncClient]):
    """Provides prompt enhancement and recommendation utilities (async).

    Accessible via ``client.prompts``.
    """

    async def stream_render(
        self,
        initial_prompt: str,
    ) -> AsyncIterator[APIEnvelope[Any]]:
        """Stream an AI-enhanced version of a prompt.

        The server progressively refines the prompt and sends intermediate
        results; the final item has ``status.code == 10000``. For a
        simpler interface that returns only the final result, use
        :meth:`render`.

        Args:
            initial_prompt: The raw prompt text to enhance.

        Yields:
            ``APIEnvelope[Any]`` envelopes as they arrive from the stream.
            Each envelope's ``.data`` dict contains a ``"prompt"`` key with
            the current version of the enhanced text.

        Example:
            >>> async for env in client.prompts.stream_render("a cat"):
            ...     if env.data:
            ...         print(env.data.get("prompt"))
        """
        async for env in self._client.request_route_stream(
            Prompts.RENDER,
            params={"prompt": initial_prompt},
            parser=self._client.parse_envelope,
        ):
            yield APIEnvelope[Any].model_validate(env)

    async def render(
        self,
        initial_prompt: str,
    ) -> str | None:
        """Enhance a prompt and return the final result.

        Convenience wrapper around :meth:`stream_render` that consumes the
        stream and returns only the final enhanced prompt string.

        Args:
            initial_prompt: The raw prompt text to enhance.

        Returns:
            The final AI-enhanced prompt string, or ``None`` if the server
            returned no data.

        Example:
            >>> enhanced = await client.prompts.render("a cat")
            >>> print(enhanced)
        """
        last: dict[str, Any] | None = None

        async for prompt in self.stream_render(
            initial_prompt,
        ):
            if prompt.data is None:
                continue
            if prompt.data and prompt.status.code != 10000:
                last = prompt.data
            elif prompt.data and prompt.status.code == 10000:
                return last.get("prompt") if last else None

        return last.get("prompt") if last else None

    async def recommend(
        self,
        model_no: str,
        base_model_name: str | None = None,
    ) -> APIEnvelope[list[PromptRecommendation]]:
        """Fetch prompt recommendations for a model.

        Args:
            model_no: Model identifier to generate recommendations for.
            base_model_name: Optional base model name to refine results.

        Returns:
            ``APIEnvelope[list[PromptRecommendation]]`` — ``.value`` is a
            list of recommendations, each with a ``.local`` text field.

        Example:
            >>> result = await client.prompts.recommend("model_abc123")
            >>> for rec in result.value or []:
            ...     print(rec.local)
        """
        payload = {"art_model_no": model_no}
        if base_model_name is not None:
            payload["base_model_name"] = base_model_name
        env = await self._client.request_route(
            Prompts.RECOMMEND,
            json=payload,
        )

        return APIEnvelope[list[PromptRecommendation]].model_validate(env)
