from __future__ import annotations

import inspect
import warnings
from contextlib import asynccontextmanager
from typing import (
    TYPE_CHECKING,
    Any,
    AsyncIterable,
    Awaitable,
    Callable,
    Iterable,
    Literal,
    Optional,
    Sequence,
    Tuple,
    Union,
    cast,
    overload,
)
from weakref import WeakValueDictionary

from htmltools import (
    HTML,
    RenderedHTML,
    Tag,
    TagAttrValue,
    TagChild,
    TagList,
    css,
)
from shiny import reactive
from shiny._deprecated import warn_deprecated
from shiny.bookmark import BookmarkState, RestoreState
from shiny.bookmark._types import BookmarkStore
from shiny.module import ResolvedId, resolve_id
from shiny.reactive._reactives import Effect_
from shiny.session import (
    get_current_session,
    require_active_session,
    session_context,
)
from shiny.types import MISSING, MISSING_TYPE, Jsonifiable, NotifyException
from shiny.ui.css import CssUnit, as_css_unit
from shiny.ui.fill import as_fill_item, as_fillable_container

from . import _utils
from ._chat_bookmark import (
    BookmarkCancelCallback,
    CancelCallback,
    ClientWithState,
    get_chatlas_state,
    is_chatlas_chat_client,
    set_chatlas_state,
)
from ._chat_normalize import message_content, message_content_chunk
from ._chat_normalize_chatlas import hide_corresponding_request, is_tool_result
from ._chat_provider_types import (
    AnthropicMessage,  # pyright: ignore[reportAttributeAccessIssue]
    GoogleMessage,
    LangChainMessage,
    OllamaMessage,
    OpenAIMessage,
    ProviderMessage,
    ProviderMessageFormat,
    as_provider_message,
)
from ._chat_tokenizer import (
    TokenEncoding,
    TokenizersEncoding,
    get_default_tokenizer,
)
from ._chat_types import (
    ChatMessage,
    ChatMessageDict,
    ClientMessage,
    TransformedMessage,
)
from ._html_deps_py_shiny import chat_deps

if TYPE_CHECKING:
    import chatlas

else:
    chatlas = object

__all__ = (
    "Chat",
    "ChatExpress",
    "chat_ui",
    "ChatMessageDict",
)


# TODO: UserInput might need to be a list of dicts if we want to support multiple
# user input content types
TransformUserInput = Callable[[str], Union[str, None]]
TransformUserInputAsync = Callable[[str], Awaitable[Union[str, None]]]
TransformAssistantResponse = Callable[[str], Union[str, HTML, None]]
TransformAssistantResponseAsync = Callable[
    [str], Awaitable[Union[str, HTML, None]]
]
TransformAssistantResponseChunk = Callable[
    [str, str, bool], Union[str, HTML, None]
]
TransformAssistantResponseChunkAsync = Callable[
    [str, str, bool], Awaitable[Union[str, HTML, None]]
]
TransformAssistantResponseFunction = Union[
    TransformAssistantResponse,
    TransformAssistantResponseAsync,
    TransformAssistantResponseChunk,
    TransformAssistantResponseChunkAsync,
]
UserSubmitFunction0 = Union[
    Callable[[], None],
    Callable[[], Awaitable[None]],
]
UserSubmitFunction1 = Union[
    Callable[[str], None],
    Callable[[str], Awaitable[None]],
]
UserSubmitFunction = Union[
    UserSubmitFunction0,
    UserSubmitFunction1,
]

ChunkOption = Literal["start", "end", True, False]

PendingMessage = Tuple[
    Any,
    ChunkOption,
    Literal["append", "replace"],
    Union[str, None],
]


class Chat:
    """
    Create a chat interface.

    A UI component for building conversational interfaces. With it, end users can submit
    messages, which will cause a `.on_user_submit()` callback to run. That callback gets
    passed the user input message, which can be used to generate a response. The
    response can then be appended to the chat using `.append_message()` or
    `.append_message_stream()`.

    Here's a rough outline for how to implement a `Chat`:

    ```python
    from shiny.express import ui

    # Create and display chat instance
    chat = ui.Chat(id="my_chat")
    chat.ui()


    # Define a callback to run when the user submits a message
    @chat.on_user_submit
    async def handle_user_input(user_input: str):
        # Create a response message stream
        response = await my_model.generate_response(user_input, stream=True)
        # Append the response into the chat
        await chat.append_message_stream(response)
    ```

    In the outline above, `my_model.generate_response()` is a placeholder for
    the function that generates a response based on the chat's messages. This function
    will look different depending on the model you're using, but it will generally
    involve passing the messages to the model and getting a response back. Also, you'll
    typically have a choice to `stream=True` the response generation, and in that case,
    you'll use `.append_message_stream()` instead of `.append_message()` to append the
    response to the chat. Streaming is preferrable when available since it allows for
    more responsive and scalable chat interfaces.

    It is also highly recommended to use a package like
    [chatlas](https://posit-dev.github.io/chatlas/) to generate responses, especially
    when responses should be aware of the chat history, support tool calls, etc.
    See this [article](https://posit-dev.github.io/chatlas/web-apps.html) to learn more.

    Parameters
    ----------
    id
        A unique identifier for the chat session. In Shiny Core, make sure this id
        matches a corresponding :func:`~shiny.ui.chat_ui` call in the UI.
    messages
        Deprecated. Use `chat.ui(messages=...)` instead.
    on_error
        How to handle errors that occur in response to user input. When `"unhandled"`,
        the app will stop running when an error occurs. Otherwise, a notification
        is displayed to the user and the app continues to run.

        * `"auto"`: Sanitize the error message if the app is set to sanitize errors,
          otherwise display the actual error message.
        * `"actual"`: Display the actual error message to the user.
        * `"sanitize"`: Sanitize the error message before displaying it to the user.
        * `"unhandled"`: Do not display any error message to the user.
    tokenizer
        Deprecated. Token counting and message trimming features will be removed in a
        future version.
    """

    def __init__(
        self,
        id: str,
        *,
        messages: Sequence[Any] = (),
        on_error: Literal["auto", "actual", "sanitize", "unhandled"] = "auto",
        tokenizer: TokenEncoding | None = None,
    ):
        if not isinstance(id, str):
            raise TypeError("`id` must be a string.")

        if messages:
            warn_deprecated(
                "`Chat(messages=...)` is deprecated. Use `.ui(messages=...)` instead."
            )

        if tokenizer is not None:
            warn_deprecated(
                "`Chat(tokenizer=...)` is deprecated. "
                "This is only relevant for `.messages(token_limits=...)` which is also now deprecated."
            )

        self.id = resolve_id(id)
        self.user_input_id = ResolvedId(f"{self.id}_user_input")
        self._transform_user: TransformUserInputAsync | None = None
        self._transform_assistant: (
            TransformAssistantResponseChunkAsync | None
        ) = None
        self._tokenizer = tokenizer

        # TODO: remove the `None` when this PR lands:
        # https://github.com/posit-dev/py-shiny/pull/793/files
        self._session = require_active_session(None)

        # Default to sanitizing until we know the app isn't sanitizing errors
        if on_error == "auto":
            on_error = "sanitize"
            app = self._session.app
            if app is not None and not app.sanitize_errors:  # type: ignore
                on_error = "actual"

        self.on_error = on_error

        # Chunked messages get accumulated (using this property) before changing state
        self._current_stream_message: str = ""
        self._current_stream_id: str | None = None
        self._pending_messages: list[PendingMessage] = []

        # For tracking message stream state when entering/exiting nested streams
        self._message_stream_checkpoint: str = ""

        # If a user input message is transformed into a response, we need to cancel
        # the next user input submit handling
        self._suspend_input_handler: bool = False

        # Keep track of effects so we can destroy them when the chat is destroyed
        self._effects: list[Effect_] = []
        self._cancel_bookmarking_callbacks: CancelCallback | None = None

        # Initialize chat state and user input effect
        with session_context(self._session):
            # Initialize message state
            self._messages: reactive.Value[tuple[TransformedMessage, ...]] = (
                reactive.Value(())
            )

            self._latest_user_input: reactive.Value[
                TransformedMessage | None
            ] = reactive.Value(None)

            @reactive.extended_task
            async def _mock_task() -> str:
                return ""

            self._latest_stream: reactive.Value[
                reactive.ExtendedTask[[], str]
            ] = reactive.Value(_mock_task)

            # TODO: deprecate messages once we start promoting managing LLM message
            # state through other means
            async def _append_init_messages():
                for msg in messages:
                    await self.append_message(msg)

            @reactive.effect
            async def _init_chat():
                await _append_init_messages()

            self._append_init_messages = _append_init_messages
            self._init_chat = _init_chat

            # When user input is submitted, transform, and store it in the chat state
            # (and make sure this runs before other effects since when the user
            #  calls `.messages()`, they should get the latest user input)
            @reactive.effect(priority=9999)
            @reactive.event(self._user_input)
            async def _on_user_input():
                msg = ChatMessage(content=self._user_input(), role="user")
                # It's possible that during the transform, a message is appended, so get
                # the length now, so we can insert the new message at the right index
                n_pre = len(self._messages())
                msg_post = await self._transform_message(msg)
                if msg_post is not None:
                    self._store_message(msg_post)
                    self._suspend_input_handler = False
                else:
                    # A transformed value of None is a special signal to suspend input
                    # handling (i.e., don't generate a response)
                    self._store_message(msg, index=n_pre)
                    await self._remove_loading_message()
                    self._suspend_input_handler = True

            self._effects.append(_init_chat)
            self._effects.append(_on_user_input)

        # Prevent repeated calls to Chat() with the same id from accumulating effects
        instance_id = self.id + "_session" + self._session.id
        instance = CHAT_INSTANCES.pop(instance_id, None)
        if instance is not None:
            instance.destroy()
        CHAT_INSTANCES[instance_id] = self

    @overload
    def on_user_submit(self, fn: UserSubmitFunction) -> Effect_: ...

    @overload
    def on_user_submit(
        self,
    ) -> Callable[[UserSubmitFunction], Effect_]: ...

    def on_user_submit(
        self, fn: UserSubmitFunction | None = None
    ) -> Effect_ | Callable[[UserSubmitFunction], Effect_]:
        """
        Define a function to invoke when user input is submitted.

        Apply this method as a decorator to a function (`fn`) that should be invoked
        when the user submits a message. This function can take an optional argument,
        which will be the user input message.

        In many cases, the implementation of `fn` should also do the following:

        1. Generate a response based on the user input.
          * If the response should be aware of chat history, use a package
             like [chatlas](https://posit-dev.github.io/chatlas/) to manage the chat
             state, or use the `.messages()` method to get the chat history.
        2. Append that response to the chat component using `.append_message()` ( or
           `.append_message_stream()` if the response is streamed).

        Parameters
        ----------
        fn
            A function to invoke when user input is submitted.

        Note
        ----
        This method creates a reactive effect that only gets invalidated when the user
        submits a message. Thus, the function `fn` can read other reactive dependencies,
        but it will only be re-invoked when the user submits a message.
        """

        def create_effect(fn: UserSubmitFunction):
            fn_params = inspect.signature(fn).parameters

            @reactive.effect
            @reactive.event(self._user_input)
            async def handle_user_input():
                if self._suspend_input_handler:
                    from shiny import req

                    req(False)
                try:
                    if len(fn_params) > 1:
                        raise ValueError(
                            "A on_user_submit function should not take more than 1 argument"
                        )
                    elif len(fn_params) == 1:
                        with warnings.catch_warnings():
                            warnings.simplefilter("ignore")
                            input = self.user_input(transform=True)
                        # The line immediately below handles the possibility of input
                        # being transformed to None. Technically, input should never be
                        # None at this point (since the handler should be suspended).
                        input = "" if input is None else input
                        afunc = _utils.wrap_async(cast(UserSubmitFunction1, fn))
                        await afunc(input)
                    else:
                        afunc = _utils.wrap_async(cast(UserSubmitFunction0, fn))
                        await afunc()
                except Exception as e:
                    await self._raise_exception(e)

            self._effects.append(handle_user_input)

            return handle_user_input

        if fn is None:
            return create_effect
        else:
            return create_effect(fn)

    async def _raise_exception(
        self,
        e: BaseException,
    ) -> None:
        if self.on_error == "unhandled":
            raise e
        else:
            await self._remove_loading_message()
            sanitize = self.on_error == "sanitize"
            msg = f"Error in Chat('{self.id}'): {str(e)}"
            raise NotifyException(msg, sanitize=sanitize) from e

    @overload
    def messages(
        self,
        *,
        format: Literal["anthropic"],
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[AnthropicMessage, ...]: ...

    @overload
    def messages(
        self,
        *,
        format: Literal["google"],
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[GoogleMessage, ...]: ...

    @overload
    def messages(
        self,
        *,
        format: Literal["langchain"],
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[LangChainMessage, ...]: ...

    @overload
    def messages(
        self,
        *,
        format: Literal["openai"],
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[OpenAIMessage, ...]: ...

    @overload
    def messages(
        self,
        *,
        format: Literal["ollama"],
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[OllamaMessage, ...]: ...

    @overload
    def messages(
        self,
        *,
        format: MISSING_TYPE = MISSING,
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[ChatMessageDict, ...]: ...

    def messages(
        self,
        *,
        format: MISSING_TYPE | ProviderMessageFormat = MISSING,
        token_limits: tuple[int, int] | None = None,
        transform_user: Literal["all", "last", "none"] = "all",
        transform_assistant: bool = False,
    ) -> tuple[ChatMessageDict | ProviderMessage, ...]:
        """
        Reactively read chat messages

        Obtain chat messages within a reactive context.

        Parameters
        ----------
        format
            Deprecated. Provider-specific message formatting will be removed in a future
            version.
        token_limits
            Deprecated. Token counting and message trimming features will be removed in
            a future version.
        transform_user
            Deprecated. Message transformation features will be removed in a future
            version.
        transform_assistant
            Deprecated. Message transformation features will be removed in a future
            version.

        Note
        ----
        Messages are listed in the order they were added. As a result, when this method
        is called in a `.on_user_submit()` callback (as it most often is), the last
        message will be the most recent one submitted by the user.

        Returns
        -------
        tuple[ChatMessage, ...]
            A tuple of chat messages.
        """

        if not isinstance(format, MISSING_TYPE):
            warn_deprecated(
                "`.messages(format=...)` is deprecated. "
                "Provider-specific message formatting will be removed in a future version. "
                "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
            )

        if token_limits is not None:
            warn_deprecated(
                "`.messages(token_limits=...)` is deprecated. "
                "Token counting and message trimming features will be removed in a future version. "
                "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
            )

        if transform_user != "all":
            warn_deprecated(
                "`.messages(transform_user=...)` is deprecated. "
                "Message transformation features will be removed in a future version. "
                "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
            )

        if transform_assistant:
            warn_deprecated(
                "`.messages(transform_assistant=...)` is deprecated. "
                "Message transformation features will be removed in a future version. "
                "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
            )

        messages = self._messages()

        # Anthropic requires a user message first and no system messages
        if format == "anthropic":
            messages = self._trim_anthropic_messages(messages)

        if token_limits is not None:
            messages = self._trim_messages(messages, token_limits, format)

        res: list[ChatMessageDict | ProviderMessage] = []
        for i, m in enumerate(messages):
            transform = False
            if m.role == "assistant":
                transform = transform_assistant
            elif m.role == "user":
                transform = transform_user == "all" or (
                    transform_user == "last" and i == len(messages) - 1
                )
            content_key = getattr(
                m, "transform_key" if transform else "pre_transform_key"
            )
            content = getattr(m, content_key)
            chat_msg = ChatMessageDict(content=str(content), role=m.role)
            if not isinstance(format, MISSING_TYPE):
                chat_msg = as_provider_message(chat_msg, format)
            res.append(chat_msg)

        return tuple(res)

    async def append_message(
        self,
        message: Any,
        *,
        icon: HTML | Tag | TagList | None = None,
    ):
        """
        Append a message to the chat.

        Parameters
        ----------
        message
            A given message can be one of the following:

            * A string, which is interpreted as markdown and rendered to HTML on the
              client.
                * To prevent interpreting as markdown, mark the string as
                  :class:`~shiny.ui.HTML`.
            * A UI element (specifically, a :class:`~shiny.ui.TagChild`).
                * This includes :class:`~shiny.ui.TagList`, which take UI elements
                  (including strings) as children. In this case, strings are still
                  interpreted as markdown as long as they're not inside HTML.
            * A dictionary with `content` and `role` keys. The `content` key can contain
              content as described above, and the `role` key can be "assistant" or
              "user".

            **NOTE:** content may include specially formatted **input suggestion** links
            (see note below).
        icon
            An optional icon to display next to the message, currently only used for
            assistant messages. The icon can be any HTML element (e.g., an
            :func:`~shiny.ui.img` tag) or a string of HTML.

        Note
        ----
        :::{.callout-note title="Input suggestions"}
        Input suggestions are special links that send text to the user input box when
        clicked (or accessed via keyboard). They can be created in the following ways:

        * `<span class='suggestion'>Suggestion text</span>`: An inline text link that
            places 'Suggestion text' in the user input box when clicked.
        * `<img data-suggestion='Suggestion text' src='image.jpg'>`: An image link with
            the same functionality as above.
        * `<span data-suggestion='Suggestion text'>Actual text</span>`: An inline text
            link that places 'Suggestion text' in the user input box when clicked.

        A suggestion can also be submitted automatically by doing one of the following:

        * Adding a `submit` CSS class or a `data-suggestion-submit="true"` attribute to
          the suggestion element.
        * Holding the `Ctrl/Cmd` key while clicking the suggestion link.

        Note that a user may also opt-out of submitting a suggestion by holding the
        `Alt/Option` key while clicking the suggestion link.
        :::

        :::{.callout-note title="Streamed messages"}
        Use `.append_message_stream()` instead of this method when `stream=True` (or
        similar) is specified in model's completion method.
        :::
        """
        # If we're in a stream, queue the message
        if self._current_stream_id:
            self._pending_messages.append((message, False, "append", None))
            return

        msg = message_content(message)
        msg = await self._transform_message(msg)
        if msg is None:
            return
        self._store_message(msg)
        await self._send_append_message(
            message=msg,
            chunk=False,
            icon=icon,
        )

    @asynccontextmanager
    async def message_stream_context(self):
        """
        Message stream context manager.

        A context manager for appending streaming messages into the chat. This context
        manager can:

        1. Be used in isolation to append a new streaming message to the chat.
            * Compared to `.append_message_stream()` this method is more flexible but
              isn't non-blocking by default (i.e., it doesn't launch an extended task).
        2. Be nested within itself
            * Nesting is primarily useful for making checkpoints to `.clear()` back
              to (see the example below).
        3. Be used from within a `.append_message_stream()`
            * Useful for inserting additional content from another context into the
              stream (e.g., see the note about tool calls below).

        Yields
        ------
        :
            A `MessageStream` class instance, which has a method for `.append()`ing
            message content chunks to as well as way to `.clear()` the stream back to
            it's initial state. Note that `.append()` supports the same message content
            types as `.append_message()`.

        Example
        -------
        ```python
        import asyncio

        from shiny import reactive
        from shiny.express import ui

        chat = ui.Chat(id="my_chat")
        chat.ui()


        @reactive.effect
        async def _():
            async with chat.message_stream_context() as msg:
                await msg.append("Starting stream...\n\nProgress:")
                async with chat.message_stream_context() as progress:
                    for x in [0, 50, 100]:
                        await progress.append(f" {x}%")
                        await asyncio.sleep(1)
                        await progress.clear()
                await msg.clear()
                await msg.append("Completed stream")
        ```

        Note
        ----
        A useful pattern for displaying tool calls in a chatbot is for the tool to
        display using `.message_stream_context()` while the the response generation is
        happening through `.append_message_stream()`. This allows the tool to display
        things like progress updates (or other "ephemeral" content) and optionally
        `.clear()` the stream back to it's initial state when ready to display the
        "final" content.
        """
        # Checkpoint the current stream state so operation="replace"  can return to it
        old_checkpoint = self._message_stream_checkpoint
        self._message_stream_checkpoint = self._current_stream_message

        # No stream currently exists, start one
        stream_id = self._current_stream_id
        is_root_stream = stream_id is None
        if is_root_stream:
            stream_id = _utils.private_random_id()
            await self._append_message_chunk(
                "", chunk="start", stream_id=stream_id
            )

        try:
            yield MessageStream(self, stream_id)
        finally:
            # Restore the checkpoint
            self._message_stream_checkpoint = old_checkpoint

            # If this was the root stream, end it
            if is_root_stream:
                await self._append_message_chunk(
                    "",
                    chunk="end",
                    stream_id=stream_id,
                )

    async def _append_message_chunk(
        self,
        message: Any,
        *,
        chunk: Literal[True, "start", "end"] = True,
        stream_id: str,
        operation: Literal["append", "replace"] = "append",
        icon: HTML | Tag | TagList | None = None,
    ) -> None:
        # If currently we're in a *different* stream, queue the message chunk
        if self._current_stream_id and self._current_stream_id != stream_id:
            self._pending_messages.append(
                (message, chunk, operation, stream_id)
            )
            return

        self._current_stream_id = stream_id

        # Normalize various message types into a ChatMessage()
        msg = message_content_chunk(message)

        if is_tool_result(message):
            await hide_corresponding_request(message)

        if operation == "replace":
            self._current_stream_message = (
                self._message_stream_checkpoint + msg.content
            )
            msg.content = self._current_stream_message
        else:
            self._current_stream_message += msg.content

        try:
            if self._needs_transform(msg):
                # Transforming may change the meaning of msg.content to be a *replace*
                # not *append*. So, update msg.content and the operation accordingly.
                chunk_content = msg.content
                msg.content = self._current_stream_message
                operation = "replace"
                msg = await self._transform_message(
                    msg, chunk=chunk, chunk_content=chunk_content
                )
                # Act like nothing happened if transformed to None
                if msg is None:
                    return
                if chunk == "end":
                    self._store_message(msg)
            elif chunk == "end":
                # When `operation="append"`, msg.content is just a chunk, but we must
                # store the full message
                self._store_message(
                    ChatMessage(
                        content=self._current_stream_message, role=msg.role
                    )
                )

            # Send the message to the client
            await self._send_append_message(
                message=msg,
                chunk=chunk,
                operation=operation,
                icon=icon,
            )
        finally:
            if chunk == "end":
                self._current_stream_id = None
                self._current_stream_message = ""
                self._message_stream_checkpoint = ""

    async def append_message_stream(
        self,
        message: Iterable[Any] | AsyncIterable[Any],
        *,
        icon: HTML | Tag | None = None,
    ):
        """
        Append a message as a stream of message chunks.

        Parameters
        ----------
        message
            An (async) iterable of message chunks. Each chunk can be one of the
            following:

            * A string, which is interpreted as markdown and rendered to HTML on the
              client.
                * To prevent interpreting as markdown, mark the string as
                  :class:`~shiny.ui.HTML`.
            * A UI element (specifically, a :class:`~shiny.ui.TagChild`).
                * This includes :class:`~shiny.ui.TagList`, which take UI elements
                  (including strings) as children. In this case, strings are still
                  interpreted as markdown as long as they're not inside HTML.
            * A dictionary with `content` and `role` keys. The `content` key can contain
              content as described above, and the `role` key can be "assistant" or
              "user".

            **NOTE:** content may include specially formatted **input suggestion** links
            (see note below).
        icon
            An optional icon to display next to the message, currently only used for
            assistant messages. The icon can be any HTML element (e.g., an
            :func:`~shiny.ui.img` tag) or a string of HTML.

        Note
        ----
        ```{.callout-note title="Input suggestions"}
        Input suggestions are special links that send text to the user input box when
        clicked (or accessed via keyboard). They can be created in the following ways:

        * `<span class='suggestion'>Suggestion text</span>`: An inline text link that
            places 'Suggestion text' in the user input box when clicked.
        * `<img data-suggestion='Suggestion text' src='image.jpg'>`: An image link with
            the same functionality as above.
        * `<span data-suggestion='Suggestion text'>Actual text</span>`: An inline text
            link that places 'Suggestion text' in the user input box when clicked.

        A suggestion can also be submitted automatically by doing one of the following:

        * Adding a `submit` CSS class or a `data-suggestion-submit="true"` attribute to
          the suggestion element.
        * Holding the `Ctrl/Cmd` key while clicking the suggestion link.

        Note that a user may also opt-out of submitting a suggestion by holding the
        `Alt/Option` key while clicking the suggestion link.
        ```

        ```{.callout-note title="Streamed messages"}
        Use this method (over `.append_message()`) when `stream=True` (or similar) is
        specified in model's completion method.
        ```

        Returns
        -------
        :
            An extended task that represents the streaming task. The `.result()` method
            of the task can be called in a reactive context to get the final state of the
            stream.
        """

        message = _utils.wrap_async_iterable(message)

        # Run the stream in the background to get non-blocking behavior
        @reactive.extended_task
        async def _stream_task():
            return await self._append_message_stream(message, icon=icon)

        _stream_task()

        self._latest_stream.set(_stream_task)

        # Since the task runs in the background (outside/beyond the current context,
        # if any), we need to manually raise any exceptions that occur
        @reactive.effect
        async def _handle_error():
            e = _stream_task.error()
            if e:
                await self._raise_exception(e)
            _handle_error.destroy()  # type: ignore

        return _stream_task

    @property
    def latest_message_stream(self) -> reactive.ExtendedTask[[], str]:
        """
        React to changes in the latest message stream.

        Reactively reads for the :class:`~shiny.reactive.ExtendedTask` behind an
        `.append_message_stream()`.

        From the return value (i.e., the extended task), you can then:

        1. Reactively read for the final `.result()`.
        2. `.cancel()` the stream.
        3. Check the `.status()` of the stream.

        Returns
        -------
        :
            An extended task that represents the streaming task. The `.result()` method
            of the task can be called in a reactive context to get the final state of the
            stream.

        Note
        ----
        If no stream has yet been started when this method is called, then it returns an
        extended task with `.status()` of `"initial"` and that it status doesn't change
        state until a message is streamed.
        """
        return self._latest_stream()

    async def _append_message_stream(
        self,
        message: AsyncIterable[Any],
        icon: HTML | Tag | None = None,
    ):
        id = _utils.private_random_id()

        empty = ChatMessageDict(content="", role="assistant")
        await self._append_message_chunk(
            empty, chunk="start", stream_id=id, icon=icon
        )

        try:
            async for msg in message:
                await self._append_message_chunk(msg, chunk=True, stream_id=id)
            return self._current_stream_message
        finally:
            await self._append_message_chunk(empty, chunk="end", stream_id=id)
            await self._flush_pending_messages()

    async def _flush_pending_messages(self):
        pending = self._pending_messages
        self._pending_messages = []
        for msg, chunk, operation, stream_id in pending:
            if chunk is False:
                await self.append_message(msg)
            else:
                await self._append_message_chunk(
                    msg,
                    chunk=chunk,
                    operation=operation,
                    stream_id=cast(str, stream_id),
                )

    # Send a message to the UI
    async def _send_append_message(
        self,
        message: TransformedMessage | ChatMessage,
        chunk: ChunkOption = False,
        operation: Literal["append", "replace"] = "append",
        icon: HTML | Tag | TagList | None = None,
    ):
        if not isinstance(message, TransformedMessage):
            message = TransformedMessage.from_chat_message(message)

        if message.role == "system":
            # System messages are not displayed in the UI
            return

        if chunk:
            msg_type = "shiny-chat-append-message-chunk"
        else:
            msg_type = "shiny-chat-append-message"

        chunk_type = None
        if chunk == "start":
            chunk_type = "message_start"
        elif chunk == "end":
            chunk_type = "message_end"

        content = message.content_client
        content_type = "html" if isinstance(content, HTML) else "markdown"

        # TODO: pass along dependencies for both content and icon (if any)
        msg = ClientMessage(
            content=str(content),
            role=message.role,
            content_type=content_type,
            chunk_type=chunk_type,
            operation=operation,
        )

        if icon is not None:
            msg["icon"] = str(icon)

        deps = message.html_deps
        if deps:
            msg["html_deps"] = deps

        # print(msg)

        await self._send_custom_message(msg_type, msg)
        # TODO: Joe said it's a good idea to yield here, but I'm not sure why?
        # await asyncio.sleep(0)

    @overload
    def transform_user_input(
        self, fn: TransformUserInput | TransformUserInputAsync
    ) -> None: ...

    @overload
    def transform_user_input(
        self,
    ) -> Callable[[TransformUserInput | TransformUserInputAsync], None]: ...

    def transform_user_input(
        self, fn: TransformUserInput | TransformUserInputAsync | None = None
    ) -> None | Callable[[TransformUserInput | TransformUserInputAsync], None]:
        """
        Deprecated. User input transformation features will be removed in a future version.
        """

        warn_deprecated(
            "The `.transform_user_input` decorator is deprecated. "
            "User input transformation features will be removed in a future version. "
            "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
        )

        def _set_transform(fn: TransformUserInput | TransformUserInputAsync):
            self._transform_user = _utils.wrap_async(fn)

        if fn is None:
            return _set_transform
        else:
            return _set_transform(fn)

    @overload
    def transform_assistant_response(
        self, fn: TransformAssistantResponseFunction
    ) -> None: ...

    @overload
    def transform_assistant_response(
        self,
    ) -> Callable[[TransformAssistantResponseFunction], None]: ...

    def transform_assistant_response(
        self,
        fn: TransformAssistantResponseFunction | None = None,
    ) -> None | Callable[[TransformAssistantResponseFunction], None]:
        """
        Deprecated. Assistant response transformation features will be removed in a future version.
        """

        warn_deprecated(
            "The `.transform_assistant_response` decorator is deprecated. "
            "Assistant response transformation features will be removed in a future version. "
            "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
        )

        def _set_transform(
            fn: TransformAssistantResponseFunction,
        ):
            nparams = len(inspect.signature(fn).parameters)
            if nparams == 1:
                fn = cast(
                    Union[
                        TransformAssistantResponse,
                        TransformAssistantResponseAsync,
                    ],
                    fn,
                )
                fn = _utils.wrap_async(fn)

                async def _transform_wrapper(
                    content: str, chunk: str, done: bool
                ):
                    return await fn(content)

                self._transform_assistant = _transform_wrapper

            elif nparams == 3:
                fn = cast(
                    Union[
                        TransformAssistantResponseChunk,
                        TransformAssistantResponseChunkAsync,
                    ],
                    fn,
                )
                self._transform_assistant = _utils.wrap_async(fn)
            else:
                raise Exception(
                    "A @transform_assistant_response function must take 1 or 3 arguments"
                )

        if fn is None:
            return _set_transform
        else:
            return _set_transform(fn)

    async def _transform_message(
        self,
        message: ChatMessage,
        chunk: ChunkOption = False,
        chunk_content: str = "",
    ) -> TransformedMessage | None:
        res = TransformedMessage.from_chat_message(message)

        if message.role == "user" and self._transform_user is not None:
            content = await self._transform_user(message.content)
        elif (
            message.role == "assistant"
            and self._transform_assistant is not None
        ):
            content = await self._transform_assistant(
                message.content,
                chunk_content,
                chunk == "end" or chunk is False,
            )
        else:
            return res

        if content is None:
            return None

        setattr(res, res.transform_key, content)
        return res

    def _needs_transform(self, message: ChatMessage) -> bool:
        if message.role == "user" and self._transform_user is not None:
            return True
        elif (
            message.role == "assistant"
            and self._transform_assistant is not None
        ):
            return True
        return False

    # Just before storing, handle chunk msg type and calculate tokens
    def _store_message(
        self,
        message: TransformedMessage | ChatMessage,
        index: int | None = None,
    ) -> None:
        if not isinstance(message, TransformedMessage):
            message = TransformedMessage.from_chat_message(message)

        with reactive.isolate():
            messages = self._messages()

        if index is None:
            index = len(messages)

        messages = list(messages)
        messages.insert(index, message)

        self._messages.set(tuple(messages))
        if message.role == "user":
            self._latest_user_input.set(message)

        return None

    def _trim_messages(
        self,
        messages: tuple[TransformedMessage, ...],
        token_limits: tuple[int, int],
        format: MISSING_TYPE | ProviderMessageFormat,
    ) -> tuple[TransformedMessage, ...]:
        n_total, n_reserve = token_limits
        if n_total <= n_reserve:
            raise ValueError(
                f"Invalid token limits: {token_limits}. The 1st value must be greater "
                "than the 2nd value."
            )

        # Since don't trim system messages, 1st obtain their total token count
        # (so we can determine how many non-system messages can fit)
        n_system_tokens: int = 0
        n_system_messages: int = 0
        n_other_messages: int = 0
        token_counts: list[int] = []
        for m in messages:
            count = self._get_token_count(m.content_server)
            token_counts.append(count)
            if m.role == "system":
                n_system_tokens += count
                n_system_messages += 1
            else:
                n_other_messages += 1

        remaining_non_system_tokens = n_total - n_reserve - n_system_tokens

        if remaining_non_system_tokens <= 0:
            raise ValueError(
                f"System messages exceed `.messages(token_limits={token_limits})`. "
                "Consider increasing the 1st value of `token_limit` or setting it to "
                "`token_limit=None` to disable token limits."
            )

        # Now, iterate through the messages in reverse order and appending
        # until we run out of tokens
        messages2: list[TransformedMessage] = []
        n_other_messages2: int = 0
        token_counts.reverse()
        for i, m in enumerate(reversed(messages)):
            if m.role == "system":
                messages2.append(m)
                continue
            remaining_non_system_tokens -= token_counts[i]
            if remaining_non_system_tokens >= 0:
                messages2.append(m)
                n_other_messages2 += 1

        messages2.reverse()

        if len(messages2) == n_system_messages and n_other_messages2 > 0:
            raise ValueError(
                f"Only system messages fit within `.messages(token_limits={token_limits})`. "
                "Consider increasing the 1st value of `token_limit` or setting it to "
                "`token_limit=None` to disable token limits."
            )

        return tuple(messages2)

    def _trim_anthropic_messages(
        self,
        messages: tuple[TransformedMessage, ...],
    ) -> tuple[TransformedMessage, ...]:
        if any(m.role == "system" for m in messages):
            raise ValueError(
                "Anthropic requires a system prompt to be specified in it's `.create()` method "
                "(not in the chat messages with `role: system`)."
            )
        for i, m in enumerate(messages):
            if m.role == "user":
                return messages[i:]

        return ()

    def _get_token_count(
        self,
        content: str,
    ) -> int:
        if self._tokenizer is None:
            self._tokenizer = get_default_tokenizer()

        encoded = self._tokenizer.encode(content)
        if isinstance(encoded, TokenizersEncoding):
            return len(encoded.ids)
        else:
            return len(encoded)

    def user_input(self, transform: bool = False) -> str | None:
        """
        Reactively read the user's message.

        Parameters
        ----------
        transform
            Whether to apply the user input transformation function (if one was
            provided).

        Returns
        -------
        str | None
            The user input message (before any transformation).

        Note
        ----
        Most users shouldn't need to use this method directly since the last item in
        `.messages()` contains the most recent user input. It can be useful for:

          1. Taking a reactive dependency on the user's input outside of a `.on_user_submit()` callback.
          2. Maintaining message state separately from `.messages()`.

        """

        if transform:
            warn_deprecated(
                "`.user_input(transform=...)` is deprecated. "
                "User input transformation features will be removed in a future version. "
                "See here for more details: https://github.com/posit-dev/shinychat/pull/91"
            )

        msg = self._latest_user_input()
        if msg is None:
            return None
        key = "content_server" if transform else "content_client"
        val = getattr(msg, key)
        return str(val)

    def _user_input(self) -> str:
        id = self.user_input_id
        return cast(str, self._session.input[id]())

    def update_user_input(
        self,
        *,
        value: str | None = None,
        placeholder: str | None = None,
        submit: bool = False,
        focus: bool = False,
    ):
        """
        Update the user input.

        Parameters
        ----------
        value
            The value to set the user input to.
        placeholder
            The placeholder text for the user input.
        submit
            Whether to automatically submit the text for the user. Requires `value`.
        focus
            Whether to move focus to the input element. Requires `value`.
        """

        if value is None and (submit or focus):
            raise ValueError(
                "An input `value` must be provided when `submit` or `focus` are `True`."
            )

        obj = _utils.drop_none(
            {
                "value": value,
                "placeholder": placeholder,
                "submit": submit,
                "focus": focus,
            }
        )

        msg = {
            "id": self.id,
            "handler": "shiny-chat-update-user-input",
            "obj": obj,
        }

        self._session._send_message_sync({"custom": {"shinyChatMessage": msg}})

    def set_user_message(self, value: str):
        """
        Deprecated. Use `update_user_input(value=value)` instead.
        """

        warn_deprecated(
            "set_user_message() is deprecated. Use update_user_input(value=value) instead."
        )

        self.update_user_input(value=value)

    async def clear_messages(self):
        """
        Clear all chat messages.
        """
        self._messages.set(())
        await self._send_custom_message("shiny-chat-clear-messages", None)

    def destroy(self):
        """
        Destroy the chat instance.
        """
        self._destroy_effects()
        self._destroy_bookmarking()

    def _destroy_effects(self):
        for x in self._effects:
            x.destroy()
        self._effects.clear()

    def _destroy_bookmarking(self):
        if not self._cancel_bookmarking_callbacks:
            return

        self._cancel_bookmarking_callbacks()
        self._cancel_bookmarking_callbacks = None

    async def _remove_loading_message(self):
        await self._send_custom_message(
            "shiny-chat-remove-loading-message", None
        )

    async def _send_custom_message(
        self, handler: str, obj: ClientMessage | None
    ):
        await self._session.send_custom_message(
            "shinyChatMessage",
            {
                "id": self.id,
                "handler": handler,
                "obj": obj,
            },
        )

    def enable_bookmarking(
        self,
        client: ClientWithState | chatlas.Chat[Any, Any],
        /,
        *,
        bookmark_on: Optional[Literal["response"]] = "response",
    ) -> CancelCallback:
        """
        Enable bookmarking for the chat instance.

        This method registers `on_bookmark` and `on_restore` hooks on `session.bookmark`
        (:class:`shiny.bookmark.Bookmark`) to save/restore chat state on both the `Chat`
        and `client=` instances. In order for this method to actually work correctly, a
        `bookmark_store=` must be specified in `shiny.App()`.

        Parameters
        ----------
        client
            The chat client instance to use for bookmarking. This can be a Chat model
            provider from [chatlas](https://posit-dev.github.io/chatlas/), or more
            generally, an instance following the `ClientWithState` protocol.
        bookmark_on
            The event to trigger the bookmarking on. Supported values include:

            - `"response"` (the default): a bookmark is triggered when the assistant is done responding.
            - `None`: no bookmark is triggered

            When this method triggers a bookmark, it also updates the URL query string to reflect the bookmarked state.


        Raises
        ------
        ValueError
            If the Shiny App does have bookmarking enabled.

        Returns
        -------
        :
            A callback to cancel the bookmarking hooks.
        """

        session = get_current_session()
        if session is None or session.is_stub_session():
            return BookmarkCancelCallback(lambda: None)

        if session.bookmark.store == "disable":
            raise ValueError(
                "Bookmarking requires a `bookmark_store` to be set. "
                "Please set `bookmark_store=` in `shiny.App()` or `shiny.express.app_opts()."
            )

        resolved_bookmark_id_str = str(self.id)
        resolved_bookmark_id_msgs_str = resolved_bookmark_id_str + "--msgs"
        get_state: Callable[[], Awaitable[Jsonifiable]]
        set_state: Callable[[Jsonifiable], Awaitable[None]]

        # Retrieve get_state/set_state functions from the client
        if isinstance(client, ClientWithState):
            # Do client with state stuff here
            get_state = _utils.wrap_async(client.get_state)
            set_state = _utils.wrap_async(client.set_state)

        elif is_chatlas_chat_client(client):
            get_state = get_chatlas_state(client)
            set_state = set_chatlas_state(client)

        else:
            raise ValueError(
                "Bookmarking requires a client that supports "
                "`async def get_state(self) -> shiny.types.Jsonifiable` (which returns an object that can be used when bookmarking to save the state of the `client=`) and "
                "`async def set_state(self, value: Jsonifiable)` (which should restore the `client=`'s state given the `state=`)."
            )

        # Reset prior bookmarking hooks
        self._destroy_bookmarking()

        # Must use `root_session` as the id is already resolved. :-/
        # Using a proxy session would double-encode the proxy-prefix
        root_session = session.root_scope()
        root_session.bookmark.exclude.append(self.id + "_user_input")

        # ###########
        # Bookmarking

        if bookmark_on is not None:
            # When ever the bookmark is requested, update the query string (indep of store type)
            @root_session.bookmark.on_bookmarked
            async def _(url: str):
                await session.bookmark.update_query_string(url)

        if bookmark_on == "response":

            @reactive.effect
            @reactive.event(
                lambda: self.messages(format=MISSING), ignore_init=True
            )
            async def _():
                messages = self.messages(format=MISSING)

                if len(messages) == 0:
                    return

                last_message = messages[-1]

                if last_message.get("role") == "assistant":
                    await session.bookmark()

        ###############
        # Client Bookmarking

        @root_session.bookmark.on_bookmark
        async def _on_bookmark_client(state: BookmarkState):
            if resolved_bookmark_id_str in state.values:
                raise ValueError(
                    f'Bookmark value with id (`"{resolved_bookmark_id_str}"`) already exists.'
                )

            with reactive.isolate():
                state.values[resolved_bookmark_id_str] = await get_state()

        @root_session.bookmark.on_restore
        async def _on_restore_client(state: RestoreState):
            if resolved_bookmark_id_str not in state.values:
                return

            # Retrieve the chat turns from the bookmark state
            info = state.values[resolved_bookmark_id_str]
            await set_state(info)

        ###############
        # UI Bookmarking

        @root_session.bookmark.on_bookmark
        def _on_bookmark_ui(state: BookmarkState):
            if resolved_bookmark_id_msgs_str in state.values:
                raise ValueError(
                    f'Bookmark value with id (`"{resolved_bookmark_id_msgs_str}"`) already exists.'
                )

            with reactive.isolate():
                # This does NOT contain the `chat.ui(messages=)` values.
                # When restoring, the `chat.ui(messages=)` values will need to be kept
                # and the `ui.Chat(messages=)` values will need to be reset
                state.values[resolved_bookmark_id_msgs_str] = self.messages(
                    format=MISSING
                )

        # Attempt to stop the initialization of the `ui.Chat(messages=)` messages
        self._init_chat.destroy()

        @root_session.bookmark.on_restore
        async def _on_restore_ui(state: RestoreState):
            # Do not call `self.clear_messages()` as it will clear the
            # `chat.ui(messages=)` in addition to the `self.messages()`
            # (which is not what we want).

            # We always want to keep the `chat.ui(messages=)` values
            # and `self.messages()` are never initialized due to
            # calling `self._init_chat.destroy()` above

            if resolved_bookmark_id_msgs_str not in state.values:
                # If no messages to restore, display the `__init__(messages=)` messages
                await self._append_init_messages()
                return

            msgs: list[Any] = state.values[resolved_bookmark_id_msgs_str]
            if not isinstance(msgs, list):
                raise ValueError(
                    f"Bookmark value with id (`{resolved_bookmark_id_msgs_str}`) must be a list of messages."
                )

            for message_dict in msgs:
                await self.append_message(message_dict)

        def _cancel_bookmarking():
            _on_bookmark_client()
            _on_bookmark_ui()
            _on_restore_client()
            _on_restore_ui()

        # Store the callbacks to be able to destroy them later
        self._cancel_bookmarking_callbacks = _cancel_bookmarking

        return BookmarkCancelCallback(_cancel_bookmarking)


class ChatExpress(Chat):
    def ui(
        self,
        *,
        messages: Optional[Sequence[TagChild | ChatMessageDict]] = None,
        placeholder: str = "Enter a message...",
        width: CssUnit = "min(680px, 100%)",
        height: CssUnit = "auto",
        fill: bool = True,
        icon_assistant: HTML | Tag | TagList | None = None,
        **kwargs: TagAttrValue,
    ) -> Tag:
        """
        Create a UI element for this `Chat`.

        Parameters
        ----------
        messages
            A sequence of messages to display in the chat. Each message can be either a
            string or a dictionary with `content` and `role` keys. The `content` key
            should contain the message text, and the `role` key can be "assistant" or
            "user".
        placeholder
            Placeholder text for the chat input.
        width
            The width of the UI element.
        height
            The height of the UI element.
        fill
            Whether the chat should vertically take available space inside a fillable
            container.
        icon_assistant
            The icon to use for the assistant chat messages. Can be a HTML or a tag in
            the form of :class:`~htmltools.HTML` or :class:`~htmltools.Tag`. If `None`,
            a default robot icon is used.
        kwargs
            Additional attributes for the chat container element.
        """

        return chat_ui(
            id=self.id,
            messages=messages,
            placeholder=placeholder,
            width=width,
            height=height,
            fill=fill,
            icon_assistant=icon_assistant,
            **kwargs,
        )

    def enable_bookmarking(
        self,
        client: ClientWithState | chatlas.Chat[Any, Any],
        /,
        *,
        bookmark_store: Optional[BookmarkStore] = None,
        bookmark_on: Optional[Literal["response"]] = "response",
    ) -> CancelCallback:
        """
        Enable bookmarking for the chat instance.

        This method registers `on_bookmark` and `on_restore` hooks on `session.bookmark`
        (:class:`shiny.bookmark.Bookmark`) to save/restore chat state on both the `Chat`
        and `client=` instances. In order for this method to actually work correctly, a
        `bookmark_store=` must be specified in `shiny.express.app_opts()`.

        Parameters
        ----------
        client
            The chat client instance to use for bookmarking. This can be a Chat model
            provider from [chatlas](https://posit-dev.github.io/chatlas/), or more
            generally, an instance following the `ClientWithState` protocol.
        bookmark_store
            A convenience parameter to set the `shiny.express.app_opts(bookmark_store=)`
            which is required for bookmarking (and `.enable_bookmarking()`). If `None`,
            no value will be set.
        bookmark_on
            The event to trigger the bookmarking on. Supported values include:

            - `"response"` (the default): a bookmark is triggered when the assistant is done responding.
            - `None`: no bookmark is triggered

            When this method triggers a bookmark, it also updates the URL query string to reflect the bookmarked state.

        Raises
        ------
        ValueError
            If the Shiny App does have bookmarking enabled.

        Returns
        -------
        :
            A callback to cancel the bookmarking hooks.
        """

        if bookmark_store is not None:
            from shiny.express import app_opts

            app_opts(bookmark_store=bookmark_store)

        return super().enable_bookmarking(client, bookmark_on=bookmark_on)


def chat_ui(
    id: str,
    *,
    messages: Optional[Sequence[TagChild | ChatMessageDict]] = None,
    placeholder: str = "Enter a message...",
    width: CssUnit = "min(680px, 100%)",
    height: CssUnit = "auto",
    fill: bool = True,
    icon_assistant: HTML | Tag | TagList | None = None,
    **kwargs: TagAttrValue,
) -> Tag:
    """
    UI container for a chat component (Shiny Core).

    This function is for locating a :class:`~shiny.ui.Chat` instance in a Shiny Core
    app. If you are using Shiny Express, use the :method:`~shiny.ui.Chat.ui` method
    instead.

    Parameters
    ----------
    id
        A unique identifier for the chat UI.
    messages
        A sequence of messages to display in the chat. A given message can be one of the
        following:

        * A string, which is interpreted as markdown and rendered to HTML on the client.
            * To prevent interpreting as markdown, mark the string as
              :class:`~shiny.ui.HTML`.
        * A UI element (specifically, a :class:`~shiny.ui.TagChild`).
            * This includes :class:`~shiny.ui.TagList`, which take UI elements
              (including strings) as children. In this case, strings are still
              interpreted as markdown as long as they're not inside HTML.
        * A dictionary with `content` and `role` keys. The `content` key can contain a
          content as described above, and the `role` key can be "assistant" or "user".

        **NOTE:** content may include specially formatted **input suggestion** links
        (see :method:`~shiny.ui.Chat.append_message` for more info).
    placeholder
        Placeholder text for the chat input.
    width
        The width of the chat container.
    height
        The height of the chat container.
    fill
        Whether the chat should vertically take available space inside a fillable container.
    icon_assistant
            The icon to use for the assistant chat messages. Can be a HTML or a tag in
            the form of :class:`~htmltools.HTML` or :class:`~htmltools.Tag`. If `None`,
            a default robot icon is used.
    kwargs
        Additional attributes for the chat container element.
    """

    id = resolve_id(id)

    icon_attr = None
    if icon_assistant is not None:
        icon_attr = str(icon_assistant)

    icon_deps = None
    if isinstance(icon_assistant, (Tag, TagList)):
        icon_deps = icon_assistant.get_dependencies()

    message_tags: list[Tag] = []
    if messages is None:
        messages = []
    for x in messages:
        role = "assistant"
        content: TagChild = None
        if not isinstance(x, dict):
            content = x
        else:
            if "content" not in x:
                raise ValueError(
                    "Each message dictionary must have a 'content' key."
                )

            content = x["content"]
            if "role" in x:
                role = x["role"]

        # `content` is most likely a string, so avoid overhead in that case
        # (it's also important that we *don't escape HTML* here).
        if isinstance(content, str):
            ui: RenderedHTML = {"html": content, "dependencies": []}
        else:
            ui = TagList(content).render()

        message_tags.append(
            Tag(
                "shiny-chat-message",
                ui["dependencies"],
                content=ui["html"],
                icon=icon_attr,
                data_role=role,
            )
        )

    res = Tag(
        "shiny-chat-container",
        Tag("shiny-chat-messages", *message_tags),
        Tag(
            "shiny-chat-input",
            id=f"{id}_user_input",
            placeholder=placeholder,
        ),
        chat_deps(),
        icon_deps,
        {
            "style": css(
                width=as_css_unit(width),
                height=as_css_unit(height),
            )
        },
        id=id,
        placeholder=placeholder,
        fill=fill,
        # Also include icon on the parent so that when messages are dynamically added,
        # we know the default icon has changed
        icon_assistant=icon_attr,
        **kwargs,
    )

    if fill:
        res = as_fillable_container(as_fill_item(res))

    return res


class MessageStream:
    """
    An object to yield from a `.message_stream_context()` context manager.
    """

    def __init__(self, chat: Chat, stream_id: str):
        self._chat = chat
        self._stream_id = stream_id

    async def replace(self, message_chunk: Any):
        """
        Replace the content of the stream with new content.

        Parameters
        -----------
        message_chunk
            The new content to replace the current content.
        """
        await self._chat._append_message_chunk(
            message_chunk,
            operation="replace",
            stream_id=self._stream_id,
        )

    async def append(self, message_chunk: Any):
        """
        Append a message chunk to the stream.

        Parameters
        -----------
        message_chunk
            A message chunk to append to this stream
        """
        await self._chat._append_message_chunk(
            message_chunk,
            stream_id=self._stream_id,
        )


CHAT_INSTANCES: WeakValueDictionary[str, Chat] = WeakValueDictionary()
