# Import Libraries
import gradio as gr
import spaces
from dotenv import load_dotenv
from implementation.answer import answer_question

load_dotenv(override=True)


def format_context(context):
    result = "<h2 style='color: #ff7800;'>Relevant Context</h2>\n\n"
    for doc in context:
        result += f"<span style='color: #ff7800;'>Source: {doc.metadata['source']}</span>\n\n"
        result += doc.page_content + "\n\n"
    return result

@spaces.GPU
def chat(history):
    last_message = (
        "\n".join(map(str, history[-1]["content"]))
        if isinstance(history[-1]["content"], list)
        else history[-1]["content"]
    )
    prior = history[:-1]
    answer, context = answer_question(last_message, prior, use_rewrite=True)
    history.append({"role": "assistant", "content": answer})
    return history, format_context(context)


def main():
    def put_message_in_chatbot(message, history):
        return "", history + [{"role": "user", "content": message}]

    with gr.Blocks(title="PyComp: Simple Python Companion") as ui:
        gr.Markdown(
            "# 🏢 Meet PyComp: Simple Python Companion\nAsk me anything about Python!")

        with gr.Row():
            with gr.Column(scale=1):
                chatbot = gr.Chatbot(
                    label="💬 Conversation",
                    height=600,
                )
                message = gr.Textbox(
                    label="Your Question",
                    placeholder="Ask anything about Python",
                    show_label=False,
                )

            with gr.Column(scale=1):
                context_markdown = gr.Markdown(
                    label="📚 Retrieved Context",
                    value="*Retrieved context will appear here*",
                    container=True,
                    height=600,
                )

        message.submit(
            put_message_in_chatbot, inputs=[
                message, chatbot], outputs=[message, chatbot]
        ).then(chat, inputs=chatbot, outputs=[chatbot, context_markdown])

    ui.launch()


if __name__ == "__main__":
    main()