import streamlit as st
from user_management import add_user, authenticate_user, load_settings, save_settings, generate_prompt, initialize_database

# Initialize the database
initialize_database()

def main():
    # State management
    if "authenticated" not in st.session_state:
        st.session_state.authenticated = False
        st.session_state.username = None
        st.session_state.page = "Login"

    def login_page():
        st.title("Login")
        username = st.text_input("Username", key="login_username")
        password = st.text_input("Password", type="password", key="login_password")
        if st.button("Login"):
            if authenticate_user(username, password):
                st.success(f"Welcome back, {username}!")
                st.session_state.authenticated = True
                st.session_state.username = username
                st.session_state.page = "Main Page"
            else:
                st.error("Invalid username or password.")
        st.write("Don't have an account? [Sign Up](#)")

    def sign_up_page():
        st.title("Sign Up")
        username = st.text_input("Username", key="signup_username")
        password = st.text_input("Password", type="password", key="signup_password")
        if st.button("Sign Up"):
            if add_user(username, password):
                st.success("Account created successfully. Please login.")
            else:
                st.error("Username already exists. Please choose another.")

    def settings_page():
        st.title("Chat with AI")
        username = st.session_state.username
        settings = load_settings(username)

        # User-specific settings form
        topic = st.text_area("Topic", value=settings.get("topic", ""), key="topic")
        instructions = st.text_area("Custom Instructions", value=settings.get("instructions", ""), key="instructions")

        if st.button("Generate Prompt"):
            settings = {"topic": topic, "instructions": instructions}
            save_settings(username, settings)            
            result = generate_prompt(topic, instructions)
            st.success(result)

    # Navigation between pages
    if st.session_state.page == "Login":
        login_page()
    elif st.session_state.page == "Sign Up":
        sign_up_page()
    elif st.session_state.page == "Main Page":
        if st.session_state.authenticated:
            settings_page()
        else:
            st.error("You must log in to access this page.")
            st.session_state.page = "Login"

    # Sidebar for navigation
    st.sidebar.title("Navigation")
    if st.session_state.authenticated:
        if st.sidebar.button("Go to Main Page"):
            st.session_state.page = "Main Page"
        if st.sidebar.button("Logout"):
            st.session_state.authenticated = False
            st.session_state.username = None
            st.session_state.page = "Login"
    else:
        st.sidebar.radio(
            "Go to",
            ["Login", "Sign Up"],
            key="sidebar_nav",
            on_change=lambda: st.session_state.update(page=st.session_state.sidebar_nav)
        )

if __name__ == "__main__":
    main()