import gradio as gr
from pipeline import run_with_chain
from my_memory_logic import memory, restatement_chain

def chat_history_fn(user_input, history):
    # Convert `history` (list of [user, ai] pairs) into memory messages
    for user_msg, ai_msg in history:
        memory.chat_memory.add_user_message(user_msg)
        memory.chat_memory.add_ai_message(ai_msg)
    
    # 1) Restate the user question with chat history
    reformulated_q = restatement_chain.run({
        "chat_history": memory.chat_memory.messages,
        "input": user_input
    })
    
    # 2) Pass the reformulated question to your pipeline
    answer = run_with_chain(reformulated_q)
    
    # 3) Update memory
    memory.chat_memory.add_user_message(user_input)
    memory.chat_memory.add_ai_message(answer)
    
    # 4) Return the new message + updated history for Gradio
    history.append((user_input, answer))
    return history, history

demo = gr.ChatInterface(
    fn=chat_history_fn,
    title="DailyWellnessAI with Memory",
    description="A chat bot that remembers context using memory + question restatement."
)
demo.launch()