import gradio as gr
from my_memory_logic import run_with_session_memory

def chat_interface_fn(message, history):
    """
    Multi-turn chat function for Gradio's ChatInterface.
    Returns strings for both message and history to match Gradio's expected format.
    """
    # Initialize history if None
    history = history or []
    
    # Get answer from the session-based memory pipeline
    try:
        answer = run_with_session_memory(message, session_id_box.value)
    except Exception as e:
        print(f"Error in run_with_session_memory: {str(e)}")
        answer = "I apologize, but I encountered an error processing your request."
    
    # Return strings directly - this is what Gradio's ChatInterface expects
    return str(answer)

# 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=chat_interface_fn,
        title="DailyWellnessAI (Session-based Memory)",
        description="Ask your questions. The session_id determines your stored memory."
    )

# Launch the Gradio interface with sharing enabled
demo.launch(share=True)