import solara
import time
import random
from typing import List
from typing_extensions import TypedDict

@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])

# 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 Page():
    title = "StreamBot"
    with solara.Head():
        solara.Title(f"{title}")
    with solara.Sidebar():
        solara.Markdown(f"#{title}")
        GithubAvatar(
            "Alonso Silva Allende",
            "alonsosilvaallende",
            "https://avatars.githubusercontent.com/u/30263736?v=4",
        )

    with solara.Column(align="center"):

        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 solara.Column(align="center"):
            with solara.lab.ChatBox(style={"position": "fixed", "overflow-y": "scroll", "scrollbar-width": "none", "-ms-overflow-style": "none", "top": "4.5rem", "bottom": "10rem", "width": "60%"}):
                for item in messages.value:
                    with solara.lab.ChatMessage(
                        user=item["role"] == "user",
                        name="StreamBot" if item["role"] == "assistant" else "User",
                        avatar_background_color="#33cccc" if item["role"] == "assistant" else "#ff991f",
                        border_radius="20px",
                    ):
                        solara.Markdown(item["content"])
            solara.lab.ChatInput(send_callback=send, style={"position": "fixed", "bottom": "3rem", "width": "60%"})