import streamlit as st
import time

# In-memory "database" for simplicity
users_db = {
    '1': {'password': '1', 'name': '1', 'age': 1, 'favorite_color': 'blue'}
}

def register_user(username, password, name, age):
    """Register a new user."""
    global users_db
    if username in users_db:
        return False
    # users_db[username] = {'password': password, 'name': name, 'age': age}
    return True

def login_user(username, password):
    """Authenticate a user."""
    if username in users_db and users_db[username]['password'] == password:
        st.session_state['logged_in'] = True
        st.session_state['user'] = users_db[username]
        print('login success')
        return True
    print('login faild')
    print(users_db)
    st.session_state['logged_in'] = 'login'
    return False

def logout_user():
    """Log out the current user."""
    st.session_state['logged_in'] = 'login'
    st.session_state.pop('user', None)
    st.session_state.pop('step1_done', None)
    st.session_state.pop('step2_done', None)
    # st.experimental_rerun()

def show_registration():
    """Display the registration form."""
    st.header("Register")
    reg_username = st.text_input("Username", key="reg_user")
    reg_password = st.text_input("Password", type="password", key="reg_pass")
    reg_name = st.text_input("Name", key="reg_name")
    reg_age = st.number_input("Age", min_value=0, max_value=120, step=1, key="reg_age")
    if st.button("Register"):
        if register_user(reg_username, reg_password, reg_name, reg_age):
            st.success("Registration successful. Please log in.")
        else:
            st.error("Username already exists. Please choose another one.")

def show_login():
    """Display the login form."""
    st.header("Login")
    username = st.text_input("Username", key="login_user")
    password = st.text_input("Password", type="password", key="login_pass")
    if st.button("Login"):
        if login_user(username, password):
            st.success(f"Welcome back, {st.session_state['user']['name']}!")

        else:
            st.error("Invalid username or password.")

def show_task():
    
    st.title("Task")

    # init session_state
    if 'step1' not in st.session_state:
        st.session_state['step1'] = None
    if 'step2' not in st.session_state:
        st.session_state['step2'] = None
    if 'step3' not in st.session_state:
        st.session_state['step3'] = None
    if 'step4' not in st.session_state:
        st.session_state['step4'] = None    

    st.header('step 1: 文案准备')
    st.text_input("请输入文案关键词", key="step1_text")
    st.checkbox("风格", key="step1_voice")
    if st.button("开始生成文案"):        
        st.session_state['step1'] = 'done'
    
    if st.session_state['step1'] == 'done':
        st.header('step 2: 语音风格')
        st.text_input("语音风格关键词", key="step2_text")
        if st.button("开始生成语音"):
            st.session_state['step2'] = 'done'
    
    if st.session_state['step2'] == 'done':
        st.header('step 3: 素材风格')
        st.text_input("请输入素材风格关键词", key="step3_text")
        if st.button("开始载入素材"):
            st.session_state['step3'] = 'done'
    
    if st.session_state['step3'] == 'done':
        st.header('step 4: 视频风格')
        st.text_input("请输入视频风格关键词", key="step4_text")
        if st.button("开始合并视频"):
            st.session_state['step4'] = 'done'
    
    if st.session_state['step4'] == 'done':
        progress_bar = st.progress(0)

        # Simulate a real-time processing task
        for i in range(100):
            progress_bar.progress(i + 1)
            time.sleep(0.05)  # Simulate real-time data streaming


    # Option to logout
    if st.button("Logout"):
        logout_user()

def main():
    """Main function to control the app flow."""
    st.title("Streamlit Application with Login and Multi-step Task")

    if 'logged_in' not in st.session_state:
        st.session_state['logged_in'] = 'login'

    if st.session_state['logged_in'] == 'login':
        show_login()
    elif st.session_state['logged_in'] == 'register':
        show_registration()
    elif st.session_state['logged_in'] == True:
        show_task()

if __name__ == "__main__":
    main()