import gradio as gr
import time


from rag_agent import RAGAgent


rag_agent = RAGAgent()


class ChatBot:
    def __init__(self, rag_agent):
        self.message_history = []
        self.rag_agent = rag_agent

    def get_response(self, message):
        return self.rag_agent.get_response(message)

    def chat(self, message):
        time.sleep(1)
        bot_response = self.get_response(message)
        self.message_history.append((message, bot_response))
        return bot_response


def create_chat_interface(rag_agent=rag_agent):
    chatbot = ChatBot(rag_agent=rag_agent)

    custom_css = """
    #chatbot {
        direction: rtl;
        height: 400px;
    }
    .message {
        font-size: 16px;
        text-align: right;
    }
    .message-wrap {
        direction: rtl !important;
    }
    .message-wrap > div {
        direction: rtl !important;
        text-align: right !important;
    }
    .input-box {
        direction: rtl !important;
        text-align: right !important;
    }
    .container {
        direction: rtl;
    }
    .contain {
        direction: rtl !important;
    }
    .bubble {
        direction: rtl !important;
        text-align: right !important;
    }
    textarea, input {
        direction: rtl !important;
        text-align: right !important;
    }
    .user-message, .bot-message {
        direction: rtl !important;
        text-align: right !important;
    }
    """

    with gr.Blocks(css=custom_css) as interface:
        with gr.Column(elem_classes="container"):
            gr.Markdown("רופא שיניים אלקטרוני", rtl=True)

            chatbot_component = gr.Chatbot(
                [],
                elem_id="chatbot",
                height=400,
                rtl=True,
                elem_classes="message-wrap"
            )

            with gr.Row():
                submit_btn = gr.Button("שלח", variant="primary")
                txt = gr.Textbox(
                    show_label=False,
                    placeholder="הקלד את ההודעה שלך כאן...",
                    container=False,
                    elem_classes="input-box",
                    rtl=True
                )

            clear_btn = gr.Button("נקה צ'אט")

        def user_message(user_message, history):
            return "", history + [[user_message, None]]

        def bot_message(history):
            user_message = history[-1][0]
            bot_response = chatbot.chat(user_message)
            history[-1][1] = bot_response
            return history

        def clear_summary():
            rag_agent.conversation_summary = ""
            rag_agent.messages = []

        submit_btn.click(user_message, [txt, chatbot_component], [txt, chatbot_component], queue=False).then(
            bot_message, chatbot_component, chatbot_component
        )

        clear_btn.click(clear_summary, None, chatbot_component, queue=False)

    return interface


# Launch the interface
chat_interface = create_chat_interface(rag_agent=rag_agent)
chat_interface.launch(share=True)