import time
import gradio as gr
from tools import create_agent
from langchain_core.messages import RemoveMessage
from langchain_core.messages import trim_messages
# from toolkits import create_agent
# from langchain.schema import AIMessage, HumanMessage, SystemMessage

AGENT = create_agent()
theme = gr.themes.Default(primary_hue="red", secondary_hue="red")

def filter_msg(msg_list:list, keep_n:int) -> list:
    """Keep only last keep_n messages from chat history. Preserves structure user msg -> tool msg -> ai msg"""
    msg = trim_messages(
            msg_list,
            strategy="last",
            token_counter=len,
            max_tokens=keep_n,
            start_on="human",
            end_on=("tool", "ai"),
            include_system=True,
    )
    return [m.id for m in msg]

def agent_response(query, config, keep_n=10):
    messages = AGENT.get_state(config).values.get("messages", [])
    
    if len(messages) > keep_n:
        keep_msg_ids = filter_msg(messages, keep_n)
        AGENT.update_state(config, {"messages": [RemoveMessage(id=m.id) for m in messages if m.id not in keep_msg_ids]})
        print("msg removed")

    # Generate answer
    answer = AGENT.invoke({"messages":query}, config=config)
    return answer["messages"][-1].content


js_func = """
function refresh() {
    const url = new URL(window.location);

    if (url.searchParams.get('__theme') != 'light') {
        url.searchParams.set('__theme', 'light');
        window.location.href = url.href;
    }
}
"""


def delete_agent():
    print("del agent")
    global AGENT
    AGENT = create_agent()
    # print(AGENT.get_state(config).values.get("messages"), "\n\n")

with gr.Blocks(theme=theme, js=js_func, title="Dataltist", fill_height=True) as iface:
    gr.Markdown("# Dataltist Chatbot 🚀")
    chatbot = gr.Chatbot(show_copy_button=True, show_share_button=False, type="messages", scale=1)
    msg = gr.Textbox(lines=1, show_label=False, placeholder="Posez vos questions sur l'assurance")  #  submit_btn=True
    # clear = gr.ClearButton([msg, chatbot], value="Effacer 🗑")
    config = {"configurable": {"thread_id": "1"}} 


    def user(user_message, history: list):
        return "", history + [{"role": "user", "content": user_message}]
    
    def bot(history: list):
        bot_message = agent_response(history[-1]["content"], config)  #AGENT.invoke({"messages":history[-1]["content"]}, config=config)
        history.append({"role": "assistant", "content": ""})
        for character in bot_message:
            history[-1]['content'] += character
            # time.sleep(0.005)
            yield history

    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, chatbot, chatbot
    )
    iface.unload(delete_agent)

if __name__ == "__main__":
    iface.launch()  # share=True # auth=("admin", "admin")