"""Streamlit interface for the barber booking system."""

import streamlit as st
from c_langgraph_flow_new import run_booking_flow, SERVICE_MAPPING
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

def initialize_session():
    """Initialize the session state with default values."""
    if "messages" not in st.session_state:
        st.session_state.messages = []
        # Add initial greeting
        st.session_state.messages.append({
            "role": "assistant",
            "content": "Hello! Welcome to our barber shop. I'll help you book an appointment. Could you please tell me your name?"
        })
    
    if "booking_info" not in st.session_state:
        st.session_state.booking_info = {
            "current_node": "Greeting",
            "booking_info": {},
            "response": "",
            "messages": []
        }

def main():
    st.set_page_config(page_title="AI Barber Booking", page_icon="💇‍♂️")
    st.title("💇‍♂️ AI Barber Booking Assistant")
    
    # Initialize session
    initialize_session()
    
    # Display current state (for debugging)
    with st.sidebar:
        st.subheader("Debug Information")
        st.write("Current State:", st.session_state.booking_info.get("current_node", "Unknown"))
        st.write("\nBooking Info:", st.session_state.booking_info.get("booking_info", {}))
        if st.checkbox("Show Full Debug Info"):
            st.json(st.session_state.booking_info)
    
    # Display service cards
    st.markdown("### Our Services")
    cols = st.columns(len(SERVICE_MAPPING))
    for col, (service_name, details) in zip(cols, SERVICE_MAPPING.items()):
        with col:
            st.markdown(
                f"""
                <div style="padding: 10px; border-radius: 10px; border: 1px solid #ddd; text-align: center;">
                    <h4>{service_name.title()}</h4>
                    <p>${details['price']}<br>{details['duration']} min</p>
                </div>
                """,
                unsafe_allow_html=True
            )
    
    # Chat interface
    st.markdown("### Chat with our Booking Assistant")
    
    # Display chat history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Get user input
    if prompt := st.chat_input("Type your message here..."):
        if prompt.strip() == "":
            st.warning("Please enter a valid message.")
            return
        
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        try:
            # Update booking_info messages
            st.session_state.booking_info["messages"] = st.session_state.messages.copy()
            
            # Process user input through LangGraph flow
            with st.spinner("Processing your request..."):
                new_state = run_booking_flow(prompt, st.session_state.booking_info)
                
                # Update session state
                st.session_state.booking_info = new_state
                
                # Add assistant response to chat history
                if new_state.get("response"):
                    st.session_state.messages.append({
                        "role": "assistant",
                        "content": new_state["response"]
                    })
                    with st.chat_message("assistant"):
                        st.markdown(new_state["response"])
            
            # Force streamlit to rerun
            st.rerun()
            
        except Exception as e:
            st.error(f"An error occurred: {str(e)}")
            print(f"Error in main loop: {str(e)}")

if __name__ == "__main__":
    main()