import streamlit as st
from langchain_openai import ChatOpenAI
from langchain_utils import invoke_chain
from langchain_community.callbacks import get_openai_callback

def init_session_state():
    if "messages" not in st.session_state:
        st.session_state.messages = []
    if "connected" not in st.session_state:
        st.session_state.connected = False

def create_sidebar():
    with st.sidebar:
        st.title("Postgres Credentials")
        st.subheader("Enter your Credentials & Connect")
        
        # Database credentials
        host = st.text_input("Host", value="localhost")
        port = st.text_input("Port", value="5432")
        user = st.text_input("User", value="postgres")
        password = st.text_input("Password", type="password")
        database = st.text_input("Database")
        
        # OpenAI API key
        api_key = st.text_input("OpenAI API Key", type="password")
        
        # Connect button
        if st.button("Connect", use_container_width=True):
            try:
                # Store credentials in session state
                st.session_state.db_credentials = {
                    "host": host,
                    "port": port,
                    "user": user,
                    "password": password,
                    "database": database
                }
                st.session_state.api_key = api_key
                st.session_state.connected = True
                st.success("Successfully connected!")
            except Exception as e:
                st.error(f"Connection failed: {str(e)}")
                st.session_state.connected = False

def main():
    init_session_state()
    create_sidebar()
    
    st.title("Chat with Postgres DB")
    
    if not st.session_state.connected:
        st.info("Please enter your credentials in the sidebar and connect first.")
        return
    
    # Display the welcome message with the database icon using markdown
    st.markdown("""
        <div style='display: flex; align-items: center; gap: 10px;'>
            <span style='font-size: 24px;'>🗄️</span>
            <span>Hello! I'm a QualityKiosk's SQL assistant. Ask me anything about your database.</span>
        </div>
    """, unsafe_allow_html=True)
    
    # Display chat messages
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Chat input
    if prompt := st.chat_input("Type a message..."):
        st.session_state.messages.append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)
        
        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                with get_openai_callback() as cb:
                    response = invoke_chain(
                        prompt,
                        st.session_state.messages,
                        st.session_state.db_credentials,
                        st.session_state.api_key
                    )
                    print(f"OpenAI Stats: {cb}")
                st.markdown(response)
        
        st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    # Set dark theme and wide layout
    st.set_page_config(
        page_title="Chat with Postgres DB",
        layout="wide",
        initial_sidebar_state="expanded",
        # Optional: Add a custom theme
        menu_items={
            'Get Help': 'https://www.qualitykiosk.com',
            'About': "# Chat with Postgres DB\nA QualityKiosk's SQL Assistant"
        }
    )
    
    # Add custom CSS for dark theme and styling
    st.markdown("""
        <style>
            .stApp {
                background-color: #1E1E1E;
                color: #FFFFFF;
            }
            .stSidebar {
                background-color: #262626;
            }
            .stButton>button {
                background-color: #0E86D4;
                color: white;
            }
            .stTextInput>div>div>input {
                background-color: #333333;
                color: white;
            }
            .stMarkdown {
                color: white;
            }
        </style>
    """, unsafe_allow_html=True)
    
    main()