import gradio as gr
import uuid
from langchain_core.messages import HumanMessage
from memory_chatbot import chatbot


# -----------------------
# Generate IDs
# -----------------------
def new_thread_id():
    return "thread_" + str(uuid.uuid4())[:8]

def new_user_id():
    return "user_" + str(uuid.uuid4())[:6]


# -----------------------
# Initial User + Thread
# -----------------------
first_user = new_user_id()
first_thread = new_thread_id()


# -----------------------
# Chat Function
# -----------------------
def chat_fn(message, chat_history, thread_id, user_id, all_threads):

    # tool_logs = []

    config = {
        "configurable": {
            "thread_id": thread_id,
            "user_id": user_id
        }
    }

    res = chatbot.invoke(
        {"messages": [HumanMessage(content=message)]},
        config=config
    )

    ai_message = res["messages"][-1].content


    chat_history.append({"role": "user", "content": message})
    chat_history.append({"role": "assistant", "content": ai_message})

    all_threads[(user_id, thread_id)] = chat_history

    return "", chat_history, all_threads


# -----------------------
# Create New Thread
# -----------------------
def create_new_thread(user_id, all_threads):

    new_tid = new_thread_id()
    all_threads[(user_id, new_tid)] = []

    return (
        all_threads,
        gr.update(choices=[t for (u, t) in all_threads if u == user_id], value=new_tid),
        new_tid,
        []
    )


# -----------------------
# Switch Thread
# -----------------------
def switch_thread(thread_id, user_id, all_threads):

    history = all_threads.get((user_id, thread_id), [])
    return thread_id, history


# -----------------------
# Create New User
# -----------------------
def create_new_user(all_threads):

    new_uid = new_user_id()
    new_tid = new_thread_id()

    all_threads[(new_uid, new_tid)] = []

    return (
        all_threads,
        gr.update(choices=list(set([u for (u, _) in all_threads])), value=new_uid),
        new_uid,
        gr.update(choices=[new_tid], value=new_tid),
        new_tid,
        []
    )


# -----------------------
# Switch User
# -----------------------
def switch_user(user_id, all_threads):

    user_threads = [t for (u, t) in all_threads if u == user_id]

    if not user_threads:
        new_tid = new_thread_id()
        all_threads[(user_id, new_tid)] = []
        user_threads = [new_tid]

    first_thread = user_threads[0]
    history = all_threads.get((user_id, first_thread), [])

    return (
        user_id,
        gr.update(choices=user_threads, value=first_thread),
        first_thread,
        history
    )


# -----------------------
# UI
# -----------------------
with gr.Blocks() as demo:

    user_state = gr.State(first_user)
    thread_state = gr.State(first_thread)

    threads_state = gr.State({
        (first_user, first_thread): []
    })

    with gr.Row():

        # Sidebar
        with gr.Column(scale=1):

            gr.Markdown("## 👤 Users")

            user_dropdown = gr.Dropdown(
                choices=[first_user],
                value=first_user,
                label="Select User"
            )

            new_user_btn = gr.Button("➕ New User")

            gr.Markdown("## 🧵 Threads")

            thread_radio = gr.Radio(
                choices=[first_thread],
                value=first_thread,
                label="Select Thread"
            )

            new_chat_btn = gr.Button("➕ New Chat")

        # Chat Area
        with gr.Column(scale=4):

            chatbot_ui = gr.Chatbot()
            msg_box = gr.Textbox(placeholder="Type message and press Enter")

            msg_box.submit(
                chat_fn,
                inputs=[msg_box, chatbot_ui, thread_state, user_state, threads_state],
                outputs=[msg_box, chatbot_ui, threads_state],
            )

    # -----------------------
    # New User
    # -----------------------
    new_user_btn.click(
        create_new_user,
        inputs=threads_state,
        outputs=[
            threads_state,
            user_dropdown,
            user_state,
            thread_radio,
            thread_state,
            chatbot_ui
        ]
    )

    # -----------------------
    # Switch User
    # -----------------------
    user_dropdown.change(
        switch_user,
        inputs=[user_dropdown, threads_state],
        outputs=[
            user_state,
            thread_radio,
            thread_state,
            chatbot_ui
        ]
    )

    # -----------------------
    # New Thread
    # -----------------------
    new_chat_btn.click(
        create_new_thread,
        inputs=[user_state, threads_state],
        outputs=[
            threads_state,
            thread_radio,
            thread_state,
            chatbot_ui
        ]
    )

    # -----------------------
    # Switch Thread
    # -----------------------
    thread_radio.change(
        switch_thread,
        inputs=[thread_radio, user_state, threads_state],
        outputs=[thread_state, chatbot_ui]
    )


demo.launch()









# import gradio as gr
# import uuid
# from langchain_core.messages import HumanMessage
# from memory_chatbot import chatbot


# # -----------------------
# # Generate IDs
# # -----------------------
# def new_thread_id():
#     return "thread_" + str(uuid.uuid4())[:8]

# def new_user_id():
#     return "user_" + str(uuid.uuid4())[:6]


# # -----------------------
# # Chat Function
# # -----------------------
# def chat_fn(message, chat_history, thread_id, user_id):

#     config = {
#         "configurable": {
#             "thread_id": thread_id,
#             "user_id": user_id
#         }
#     }

#     res = chatbot.invoke(
#         {"messages": [HumanMessage(content=message)]},
#         config=config
#     )

#     ai_message = res["messages"][-1].content

#     chat_history.append({"role": "user", "content": message})
#     chat_history.append({"role": "assistant", "content": ai_message})

#     return "", chat_history


# # -----------------------
# # New Chat (Reset Thread)
# # -----------------------
# def new_chat():
#     return new_thread_id(), []


# # -----------------------
# # UI
# # -----------------------
# with gr.Blocks() as demo:

#     # Hidden states
#     user_state = gr.State(new_user_id())
#     thread_state = gr.State(new_thread_id())

#     gr.Markdown("# 🤖 AI Chatbot")

#     chatbot_ui = gr.Chatbot(height=500)
#     msg_box = gr.Textbox(placeholder="Type your message...")

#     with gr.Row():
#         send_btn = gr.Button("Send")
#         new_chat_btn = gr.Button("🧹 New Chat")

#     # Send message
#     send_btn.click(
#         chat_fn,
#         inputs=[msg_box, chatbot_ui, thread_state, user_state],
#         outputs=[msg_box, chatbot_ui],
#     )

#     msg_box.submit(
#         chat_fn,
#         inputs=[msg_box, chatbot_ui, thread_state, user_state],
#         outputs=[msg_box, chatbot_ui],
#     )

#     # New chat (reset thread)
#     new_chat_btn.click(
#         new_chat,
#         outputs=[thread_state, chatbot_ui]
#     )


# demo.launch()