import gradio as gr
from my_memory_logic import run_with_session_memory

def chat_interface_fn(message, history, session_id):
    """
    Simple chat function for Gradio's ChatInterface that only shows the most recent question and answer pair.
    The session_id is used to maintain session-based memory.
    """
    # 1) Get answer from the session-based memory pipeline
    answer = run_with_session_memory(message, session_id)

    # 2) Append only the new question and answer pair to history
    history.append((message, answer))

    # 3) Convert the most recent question-answer pair to message dictionaries for display
    message_dicts = []
    for user_msg, ai_msg in history[-1:]:  # Only use the latest Q&A
        message_dicts.append({"role": "user", "content": user_msg})
        message_dicts.append({"role": "assistant", "content": ai_msg})

    # Return the message dicts and updated history
    return message_dicts, history

# Custom CSS for chat interface
my_chat_css = """
.gradio-container {
  margin: auto;
}
.user .wrap {
  text-align: right !important;
}
.assistant .wrap {
  text-align: left !important;
}
"""

# Set up Gradio interface
with gr.Blocks(css=my_chat_css) as demo:
    gr.Markdown("### DailyWellnessAI (User on right, Assistant on left)")
    session_id_box = gr.Textbox(label="Session ID", value="abc123", interactive=True)
    
    chat_interface = gr.ChatInterface(
        fn=lambda msg, hist: chat_interface_fn(msg, hist, session_id_box.value),
        title="DailyWellnessAI (Session-based Memory)",
        description="Ask your questions. The session_id determines your stored memory."
    )

# Launch the Gradio interface
demo.launch()