import solara
import time
import random
from typing import List
from typing_extensions import TypedDict
from solara.components.input import use_change

# Streamed response emulator
def response_generator():
    response = random.choice(
        [
            "Hello! How can I assist you today?",
            "Hey there! If you have any questions or need help with something, feel free to ask.",
        ]
    )
    for word in response.split():
        yield word + " "
        time.sleep(0.05)

class MessageDict(TypedDict):
    role: str
    content: str

messages: solara.Reactive[List[MessageDict]] = solara.reactive([])

def add_chunk_to_ai_message(chunk: str):
    messages.value = [
        *messages.value[:-1],
        {
            "role": "assistant",
            "content": messages.value[-1]["content"] + chunk,
        },
    ]

@solara.component
def GithubAvatar(name: str, handle: str, img: str):
    with solara.v.Html(tag="a", attributes={"href": f"https://github.com/{handle}/", "target": "_blank"}):
        with solara.v.ListItem(class_="pa-0"):
            with solara.v.ListItemAvatar(color="grey darken-3"):
                solara.v.Img(
                    class_="elevation-6",
                    src=img,
                )
            with solara.v.ListItemContent():
                solara.v.ListItemTitle(children=["By " + name])

import uuid
from typing import Callable, Dict, List, Optional, Union
from typing_extensions import Literal
@solara.component
def ChatBox(
    children: List[solara.Element] = [],
    style: Optional[Union[str, Dict[str, str]]] = None,
    classes: List[str] = [],
):
    """
    The ChatBox component is a container for ChatMessage components.
    Its primary use is to ensure the proper ordering of messages,
    using `flex-direction: column-reverse` together with `reversed(messages)`.

    # Arguments

    * `children`: A list of child components.
    * `style`: CSS styles to apply to the component. Either a string or a dictionary.
    * `classes`: A list of CSS classes to apply to the component.
    """
    style_flat = solara.util._flatten_style(style)
    style_flat += " background-color:transparent!important;"
    if "flex-grow" not in style_flat:
        style_flat += " flex-grow: 1;"
    if "flex-direction" not in style_flat:
        style_flat += " flex-direction: column-reverse;"
    if "overflow-y" not in style_flat:
        style_flat += " overflow-y: auto;"

    #classes += ["chat-box"]
    with solara.Column(
        style=style_flat,
        classes=classes,
    ):
        for child in list(reversed(children)):
            solara.display(child, style={"background-color":"transparent!important"})

@solara.component
def ChatInput(
    send_callback: Optional[Callable] = None,
    disabled: bool = False,
    style: Optional[Union[str, Dict[str, str]]] = None,
    classes: List[str] = [],
):
    """
    The ChatInput component renders a text input together with a send button.

    # Arguments

    * `send_callback`: A callback function for when the user presses enter or clicks the send button.
    * `disabled`: Whether the input should be disabled. Useful for disabling sending further messages while a chatbot is replying,
        among other things.
    * `style`: CSS styles to apply to the component. Either a string or a dictionary. These styles are applied to the container component.
    * `classes`: A list of CSS classes to apply to the component. Also applied to the container.
    """
    message, set_message = solara.use_state("")  # type: ignore
    style_flat = solara.util._flatten_style(style)

    if "align-items" not in style_flat:
        style_flat += " align-items: center;"

    with solara.Row(style=style_flat, classes=classes):

        def send(*ignore_args):
            if message != "" and send_callback is not None:
                send_callback(message)
                set_message("")

        message_input = solara.v.TextField(
            label="Send message...",
            v_model=message,
            on_v_model=set_message,
            rounded=True,
            filled=True,
            hide_details=True,
            style_="flex-grow: 1;",
            disabled=disabled,
            color="secondary",
        )

        use_change(message_input, send, update_events=["keyup.enter"])

        button = solara.v.Btn(icon=True, children=[solara.v.Icon(children=["mdi-send"])], disabled=message == "")

        use_change(button, send, update_events=["click"])

@solara.component
def ChatMessage(
    children: Union[List[solara.Element], str],
    user: bool = False,
    avatar: Union[solara.Element, str, Literal[False], None] = None,
    name: Optional[str] = None,
    color: Optional[str] = None,
    avatar_background_color: Optional[str] = None,
    border_radius: Optional[str] = None,
    notch: bool = False,
    style: Optional[Union[str, Dict[str, str]]] = None,
    classes: List[str] = [],
):
    """
    The ChatMessage component renders a message. Messages with `user=True` are rendered on the right side of the screen,
    all others on the left.

    # Arguments

    * `children`: A list of child components.
    * `user`: Whether the message is from the current user or not.
    * `avatar`: An avatar to display next to the message. Can be a string representation of a URL or Material design icon name,
        a solara Element, False to disable avatars altogether, or None to display initials based on `name`.
    * `name`: The name of the user who sent the message.
    * `color`: The background color of the message. Defaults to `rgba(0,0,0,.06)`. Can be any valid CSS color.
    * `avatar_background_color`: The background color of the avatar. Defaults to `color` if left as `None`.
    * `border_radius`: Sets the roundness of the corners of the message. Defaults to `None`,
        which applies the default border radius of a `solara.Column`, i.e. `4px`.
    * `notch`: Whether to display a speech bubble style notch on the side of the message.
    * `style`: CSS styles to apply to the component. Either a string or a dictionary. Applied to the container of the message.
    * `classes`: A list of CSS classes to apply to the component. Applied to the same container.
    """
    style_flat = solara.util._flatten_style(style)

    if "border-radius" not in style_flat:
        style_flat += f" border-radius: {border_radius if border_radius is not None else ''};"
    if f"border-top-{'right' if user else 'left'}-radius" not in style_flat:
        style_flat += f" border-top-{'right' if user else 'left'}-radius: 0;"
    if "padding" not in style_flat:
        style_flat += " padding: .5em 1.5em;"

    msg_uuid = solara.use_memo(lambda: str(uuid.uuid4()), dependencies=[])
    with solara.Row(
        justify="end" if user else "start",
        style={"color":"transparent!important", "flex-direction": "row-reverse" if user else "row", "padding": "0px"},
    ):
        if avatar is not False:
            with solara.v.Avatar(color=avatar_background_color if avatar_background_color is not None else color):
                if avatar is None and name is not None:
                    initials = "".join([word[:1] for word in name.split(" ")])
                    solara.HTML(tag="span", unsafe_innerHTML=initials, classes=["headline"])
                elif isinstance(avatar, solara.Element):
                    solara.display(avatar)
                elif isinstance(avatar, str) and avatar.startswith("mdi-"):
                    solara.v.Icon(children=[avatar])
                else:
                    solara.HTML(tag="img", attributes={"src": avatar, "width": "100%"})
        classes_new = classes + ["chat-message-" + msg_uuid, "right" if user else "left"]
        with solara.Column(
            classes=classes_new,
            gap=0,
            style=style_flat,
        ):
            if name is not None:
                solara.Text(name, style="font-weight: bold;", classes=["message-name", "right" if user else "left"])
            for child in children:
                if isinstance(child, solara.Element):
                    solara.display(child)
                else:
                    solara.Markdown(child)
        # we use the uuid to generate 'scoped' CSS, i.e. css that only applies to the component instance.
        extra_styles = (
            f""".chat-message-{msg_uuid}:before{{
                content: '';
                position: absolute;
                width: 0;
                height: 0;
                border: 6px solid;
                top: 0;
                color: transparent;
            }}
            .chat-message-{msg_uuid}.left:before{{
                left: -12px;
                color: transparent;
                border-color: var(--color) var(--color) transparent transparent;
            }}
            .chat-message-{msg_uuid}.right:before{{
                right: -12px;
                border-color: var(--color) transparent transparent var(--color);
                color: transparent;
            }}"""
            if notch
            else ""
        )
        solara.Style(
            f"""
            .chat-message-{msg_uuid}{{
                color: transparent!important;
                max-width: 75%;
                position: relative;
            }}
            .chat-message-{msg_uuid}.left{{
                    color: transparent!important;
                    border-top-left-radius: 0;
                    background-color:var(--color);
                    { "margin-left: 10px !important;" if notch else ""}
            }}
            .chat-message-{msg_uuid}.right{{
                    color: transparent!important;
                    border-top-right-radius: 0;
                    background-color:var(--color);
                    { "margin-right: 10px !important;" if notch else ""}
            }}
            {extra_styles}
            """
        )

@solara.component
def Page():
    solara.lab.theme.themes.light.primary = "#ff0000"
    solara.lab.theme.themes.light.secondary = "#0000ff"
    solara.lab.theme.themes.dark.primary = "#ff0000"
    solara.lab.theme.themes.dark.secondary = "#0000ff"
    title = "Customized StreamBot"
    with solara.Head():
        solara.Title(f"{title}")
    with solara.AppBar():
        solara.lab.ThemeToggle(enable_auto=False)
    with solara.Sidebar():
        solara.Markdown(f"#{title}")
        GithubAvatar(
            "Alonso Silva Allende",
            "alonsosilvaallende",
            "https://avatars.githubusercontent.com/u/30263736?v=4",
        )
    with solara.Columns([1,1],style="padding: 20em 10em 20em 10em;color: transparent!important; background-color: #00ff00!important; background-image: url(https://wallpapercave.com/wp/DlGpnB5.jpg)"):
        with solara.Column(align="center",style={"background-color":"red"}):
            user_message_count = len([m for m in messages.value if m["role"] == "user"])
    
            def send(message):
                messages.value = [
                    *messages.value,
                    {"role": "user", "content": message},
                ]
    
            def response(message):
                messages.value = [*messages.value, {"role": "assistant", "content": ""}]
                for chunk in response_generator():
                    add_chunk_to_ai_message(chunk)
    
    
            def result():
                if messages.value !=[]: response(messages.value[-1]["content"])
    
            result = solara.lab.use_task(result, dependencies=[user_message_count])  # type: ignore
            
            #with ChatBox(style={"background-color":"transparent!important","flex-grow":"1","position": "fixed", "bottom": "10rem", "width": "50%"}):
            with ChatBox(style={"position": "fixed", "overflow-y": "scroll","scrollbar-width": "none", "-ms-overflow-style": "none", "top": "4.5rem", "bottom": "10rem", "width": "50%"}):
                for item in messages.value:
                    with ChatMessage(
                        user=item["role"] == "user",
                        name="StreamBot" if item["role"] == "assistant" else "User",
                        notch=True,
                        avatar="https://avatars.githubusercontent.com/u/127238744?v=4" if item["role"] == "user" else "https://avatars.githubusercontent.com/u/784313?v=4",
                        avatar_background_color="#33cccc" if item["role"] == "assistant" else "#ff991f",
                        border_radius="20px",
                        style={"color":"red!important","background-color":"orange!important","font-family":"Comic Sans MS"} if item["role"] == "user" else {"color":"red!important", "background-color":"aqua!important","font-family":"Comic Sans MS"},
                    ):
                        solara.Markdown(
                                item["content"],
                                style={"color":"green!important","backgound-color":"transpartent!important","font-family":"Comic Sans MS"} if item["role"] == "user" else {"color":"blue!important","font-family":"Comic Sans MS"}
                        )
            ChatInput(send_callback=send, style={"position": "fixed", "bottom": "3rem", "width": "60%", "color":"green!important", "background-color":"red!important"})