import json
import gradio as gr
from ai_llm import AI_LLM
from pathway import Pathway

with open("pathways/Virelia/Virelia.json", "r") as file:
    pathway_json = json.load(file)
ai_llm = AI_LLM()
temp_pathway = Pathway(pathway_json, ai_llm)

# Initial chatbot messages
initial_ui = temp_pathway.get_current_stage_inital_ui()
chatbot_messages = []
image_url = ""

if initial_ui:
    if "text" in initial_ui:
        for msg in initial_ui["text"]:
            chatbot_messages.append({
                "role": "assistant",
                "content": msg
            })
    if "image" in initial_ui:
        image_url = initial_ui["image"]

sessions = {}

def initialize_session(request: gr.Request):
    sessions[request.session_hash] = Pathway(pathway_json, ai_llm)

def close_session(request: gr.Request):
    if request.session_hash in sessions:
        del sessions[request.session_hash]

def next_turn(request:gr.Request, user_input, chatbot_history):
    if user_input.strip() == "":
        return "", chatbot_history, gr.update(), gr.update(), gr.update()
    
    pathway = sessions[request.session_hash]
    chatbot_history.append({
        "role": "user",
        "content": user_input
    })
    pathway_complete, response, next_stage_ui, checker_ui = pathway.next_step(user_input)
    for r in response["text"]:
        chatbot_history.append({
            "role": "assistant",
            "content": r
        })

    if next_stage_ui:
        if "text" in next_stage_ui:
            for msg in next_stage_ui["text"]:
                chatbot_history.append({
                    "role": "assistant",
                    "content": msg
                })

    if checker_ui:
        checker_text = ""
        checker_title = "Hint"
        if "text" in checker_ui:
            for msg in checker_ui["text"]:
                checker_text += "<span style='color:orange; font-weight:bold;'>"
                checker_text += msg + "</span><br>"
        if "title" in checker_ui:
            checker_title = checker_ui["title"]

        chatbot_history.append(gr.ChatMessage(
            role="assistant",
            content=checker_text,
            metadata={"title": "<span style='color:red; font-weight:bold;'>" + checker_title +"</span>", "status": "done"}
        ))

    if "image" in response:
        stage_image = response["image"]
    else:
        stage_image = gr.update()

    if "video" in response:
        stage_video = gr.update(value=response["video"], visible=True)
    else:
        stage_video = gr.update(visible=False)

    if pathway_complete:
        chatbot_history.append({
            "role": "assistant",
            "content": "Congratulations! You have completed the quest."
        })

        pathway_complete_ui = pathway.get_pathway_complete_ui()
        if pathway_complete_ui:
            title = "Quest Complete"
            text = ""
            if "title" in pathway_complete_ui:
                title = pathway_complete_ui["title"]
            if "text" in pathway_complete_ui:
                for msg in pathway_complete_ui["text"]:
                    text += "<span style='color:orange; font-weight:bold;'>"
                    text += msg + "</span><br>"
            chatbot_history.append(gr.ChatMessage(
                role="assistant",
                content=text,
                metadata={"title": "<span style='color:red; font-weight:bold;'>" + title + "</span>", "status": "done"}
            ))
        return gr.update(interactive=False, value=""), chatbot_history, stage_image, stage_video, gr.update(interactive=False)

    return "", chatbot_history, stage_image, stage_video, gr.update()

css = """
button[aria-label="Clear"] {
    display: none !important;
}
footer {display: none !important;}
.gr-chatbot {
    max-height: 500px; /* Set the maximum height */
    overflow-y: auto;  /* Enable vertical scrolling for overflow */
}
"""

with gr.Blocks(css=css) as chatbotDisplay:
    with gr.Row():
        with gr.Column(scale=1):
            stage_image = gr.Image(value=image_url, show_label=False, show_download_button=False)
            stage_video = gr.Video(show_label=False, show_download_button=False, visible=False, autoplay=True)
        with gr.Column(scale=2):
            chatBot = gr.Chatbot(chatbot_messages,type="messages", group_consecutive_messages=False, scale=2, show_label=False)
            input = gr.Textbox(show_label=False, placeholder="Type your message here...", lines=1)
            submit_btn = gr.Button("Submit")

    input.submit(next_turn, [input, chatBot], [input, chatBot, stage_image, stage_video, submit_btn])
    submit_btn.click(next_turn, [input, chatBot], [input, chatBot, stage_image, stage_video, submit_btn])
    chatbotDisplay.load(initialize_session)
    chatbotDisplay.close(close_session)

chatbotDisplay.launch(share=False)