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


# -----------------------
# 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):

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

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

    ai_message = get_last_ai_message(res["messages"])
    tool_logs = extract_tool_logs(res["messages"], message)

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

    all_threads[(user_id, thread_id)] = {
        "chat": chat_history,
        "tools": tool_logs
    }

    return "", chat_history, all_threads, tool_logs


# -----------------------
# Create New Thread
# -----------------------

def create_new_thread(user_id, all_threads):

    new_tid = new_thread_id()
    all_threads[(user_id, new_tid)] = {
        "chat": [],
        "tools": "_No tools used yet._"
    }

    return (
        all_threads,
        gr.update(choices=[t for (u, t) in all_threads if u == user_id], value=new_tid),
        new_tid,
        [],
        "_No tools used yet._"
    )


# -----------------------
# Switch Thread
# -----------------------

def switch_thread(thread_id, user_id, all_threads):

    data = all_threads.get(
        (user_id, thread_id),
        {"chat": [], "tools": "_No tools used yet._"}
    )

    return thread_id, data["chat"], data["tools"]

# -----------------------
# 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)] = {
        "chat": [],
        "tools": "_No tools used yet._"
    }

    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,
        [],
        "_No tools used yet._"
    )


# -----------------------
# 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)] = {
            "chat": [],
            "tools": "_No tools used yet._"
        }
        user_threads = [new_tid]

    first_thread = user_threads[0]
    data = all_threads.get(
        (user_id, first_thread),
        {"chat": [], "tools": "_No tools used yet._"}
    )

    return (
        user_id,
        gr.update(choices=user_threads, value=first_thread),
        first_thread,
        data["chat"],
        data["tools"]
    )


# -----------------------
# 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): {
        "chat": [],
        "tools": "_No tools used yet._"
        }
    })

    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=7):

            chatbot_ui = gr.Chatbot()
            msg_box = gr.Textbox(placeholder="Type message and press Enter")
            tool_panel = gr.Markdown("## 🛠 Tool Activity\n_No tools used yet._")

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

    # -----------------------
    # 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,
        tool_panel
    ]
)

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

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

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


demo.launch()