from __future__ import annotations

from collections.abc import AsyncIterator, Mapping
from typing import TYPE_CHECKING, Any, cast

from ...._endpoints import GenAgent
from ...._internal._builders.gen_agent import build_stream_gen_agent_message
from ...._internal._schemas.genagent.gen_agent_messages import (
    AgentGenerationChunk,
    AgentGenerationEnd,
    AgentGenerationResult,
    AgentGenerationStreamItem,
)
from ..._base import BaseResource

if TYPE_CHECKING:
    from ...._async_client import AsyncClient
    from ....options.gen_agent_messages import GenAgentReqOptions
else:
    AsyncClient = Any


class AsyncGenAgentMessagesResource(BaseResource[AsyncClient]):
    """Manages messages in a Gen Agent conversation (async).

    Accessible via ``client.genagent.messages``.
    """

    async def stream(
        self,
        message: str,
        conv_id: str,
        *,
        parent_msg_id: str | None = None,
        start_new_topic: bool = True,
        role: str = "user",
        options: GenAgentReqOptions | None = None,
    ) -> AsyncIterator[AgentGenerationStreamItem]:
        """Send a message to the Gen Agent and stream the response.

        Args:
            message: User message text to send.
            conv_id: ID of the existing conversation to send the message to.
            parent_msg_id: Parent message ID for threaded replies.
            start_new_topic: Whether to start a new topic in the conversation.
            role: Role of the sender. Defaults to ``"user"``.
            options: Additional request options (attached images, model
                overrides, agent mode, etc.). See ``GenAgentReqOptions``.

        Yields:
            ``AgentGenerationChunk`` while text is being generated
            (``item.event == "message"``; use ``item.text`` for the
            fragment), followed by a final ``AgentGenerationEnd``
            (``item.event == "end"``) that carries message IDs and
            other metadata.

        Example:
            >>> async for item in client.genagent.messages.stream(
            ...     "Generate a fantasy portrait", conv_id
            ... ):
            ...     if item.event == "message":
            ...         print(item.text, end="", flush=True)
        """
        body = build_stream_gen_agent_message(
            conv_id=conv_id,
            start_new_topic=start_new_topic,
            message=message,
            role=role,
            parent_msg_id=parent_msg_id,
            options=options,
        )

        async for env in self._client.request_route_stream(
            GenAgent.STREAM,
            json=body.model_dump(exclude_none=True),
            parser=self._client.parse_sse_stream,
        ):
            if env is None or env.data is None:
                continue
            if env.data == "[DONE]":
                continue

            data_any: Any = env.data

            if not isinstance(data_any, Mapping):
                continue

            payload = cast("Mapping[str, Any]", data_any)

            if "choices" in payload:
                yield AgentGenerationChunk.model_validate(env)
                continue

            if "status" in payload:
                status_any = payload.get("status")
                data_obj = payload.get("data")

                end_dict: dict[str, Any] = {"event": "end", "status": status_any}

                if isinstance(data_obj, Mapping):
                    end_dict.update(dict(cast("Mapping[str, Any]", data_obj)))

                yield AgentGenerationEnd.model_validate(end_dict)
                break

            continue

    async def send(
        self,
        message: str,
        conv_id: str,
        *,
        parent_msg_id: str | None = None,
        start_new_topic: bool = True,
        role: str = "user",
        options: GenAgentReqOptions | None = None,
    ) -> AgentGenerationResult | None:
        """Send a message to the Gen Agent and return the full response.

        Convenience wrapper around :meth:`stream` that collects all chunks
        and returns the assembled result.

        Args:
            message: User message text to send.
            conv_id: ID of the existing conversation to send the message to.
            parent_msg_id: Parent message ID for threaded replies.
            start_new_topic: Whether to start a new topic in the conversation.
            role: Role of the sender. Defaults to ``"user"``.
            options: Additional request options. See ``GenAgentReqOptions``.

        Returns:
            ``AgentGenerationResult`` with ``.text`` containing the full
            assembled response and ``.value`` being the final
            ``AgentGenerationEnd`` object (message IDs, other info).
            Returns ``None`` if no content was generated.

        Example:
            >>> result = await client.genagent.messages.send(
            ...     "Generate a fantasy portrait", conv_id
            ... )
            >>> if result:
            ...     print(result.text)
        """
        parts: list[str] = []
        env = None

        async for env in self.stream(
            conv_id=conv_id,
            message=message,
            parent_msg_id=parent_msg_id,
            start_new_topic=start_new_topic,
            role=role,
            options=options,
        ):
            if env.event == "message":
                parts.append(env.text)

        if not parts or env is None:
            return None

        if env.event == "end":
            return AgentGenerationResult(text="".join(parts), value=env)
        return AgentGenerationResult(text="".join(parts), value=AgentGenerationEnd())
