import asyncio
import threading
import time

from reactpy import component, html, hooks, use_effect, run
from reactpy.backend.starlette import configure
from starlette.applications import Starlette
from vectara_agentic.agent import AgentStatusType
from agent import initialize_agent, get_agent_config

from header import Header
from footer import Footer
from chatbox import ChatBox
from chatui import ChatUI
from utils import wait_message

@component
def App():
    # State hooks
    show_logs, set_show_logs = hooks.use_state(False)
    show_logs_button, set_show_logs_button = hooks.use_state(False)
    collapsed, set_collapsed = hooks.use_state(False) # Whether or not the footer is collapsed
    first_turn, set_first_turn = hooks.use_state(True)
    messages, set_messages = hooks.use_state([])
    message, set_message = hooks.use_state("")

    def use_agent_logger():
        agent_log_entries, set_agent_log_entries = hooks.use_state([])

        def reset_log_entries():
            set_agent_log_entries([])

        def add_log_entry(new_log_entry):
            set_agent_log_entries(lambda previous_entries: previous_entries + [new_log_entry])

        return agent_log_entries, add_log_entry, reset_log_entries

    # Logger setup
    log_entries, add_log_entry, reset_log_entries = use_agent_logger()

    def update_func(status_type: AgentStatusType, msg: str):
        if status_type != AgentStatusType.AGENT_UPDATE:
            output = f"{status_type.value} - {msg}"
            add_log_entry(output)

    # Agent setup
    cfg, _ = hooks.use_state(get_agent_config())
    agent, _ = hooks.use_state(initialize_agent(cfg, update_func))

    
    def start_over(event=None):
        set_messages([])
        set_first_turn(True)

    def display_message(new_messages):
        if first_turn:
            set_first_turn(False)
        set_messages(messages + list(new_messages))
        
    async def chat_response(user_message):
        response = await asyncio.to_thread(agent.chat, user_message)
        return response

    async def send_message_async(sent_message):
        response = await chat_response(sent_message)
        set_messages(messages[:-1])
        display_message(
            [
                {"user": "human", "message": sent_message},
                {"user": "bot", "message": response},
            ]
        )
        set_show_logs_button(True)
   
    def send_message(event=None):
        if message.strip():
            sent_message = message
            set_message("")
            set_show_logs_button(False)
            set_show_logs(False)
            reset_log_entries()
            display_message([{"user": "human", "message": sent_message}])
            display_message([{"user": "bot", "message": wait_message}])
            asyncio.create_task(send_message_async(sent_message))

    def send_example(ex_prompt):
        if ex_prompt.strip():
            sent_message = ex_prompt
            set_message("")
            reset_log_entries()
            display_message([{"user": "human", "message": sent_message}])
            display_message([{"user": "bot", "message": wait_message}])

            asyncio.create_task(send_message_async(sent_message))

    return html.div(
        {
            "style": {
                "backgroundColor": "#F4F1F4",
                "margin": "0",
                "padding": "0",
                "minHeight": "100vh",
                "display": "flex",
                "boxSizing": "border-box",
                "flexDirection": "column",
                "overflowX": "hidden",
                "position": "relative",
            }
        },
        html.head(
            html.style(
                """
                body {
                    margin: 0;
                    padding: 0;
                }
                """
            )
        ),
        Header(
            demo_name=cfg['demo_name'],
            short_description=cfg['short_description'],
            extra_info=cfg['extra_info'],
            start_over=start_over
        ),
        ChatUI(
            messages=messages,
            first_turn=first_turn,
            examples=cfg['examples'],
            send_example=send_example,
            log_entries=log_entries,
            show_logs=show_logs,
            set_show_logs=set_show_logs,
            show_logs_button=show_logs_button,
            collapsed=collapsed
        ),
        ChatBox(message, set_message, send_message, collapsed),
        Footer(collapsed, set_collapsed),
    )

app = Starlette()
configure(app, App)

# if __name__ == "__main__":
#     run(App)