import streamlit as st
import sqlite3
import hashlib

# Database setup
DB_FILE = "users.db"

def create_user_table():
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    cursor.execute("""
        CREATE TABLE IF NOT EXISTS users (
            username TEXT PRIMARY KEY,
            password TEXT
        )
    """)
    conn.commit()
    conn.close()

def add_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    try:
        cursor.execute("INSERT INTO users (username, password) VALUES (?, ?)", (username, hashed_password))
        conn.commit()
    except sqlite3.IntegrityError:
        st.error("Username already exists. Please choose a different username.")
    conn.close()

def authenticate_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    cursor = conn.cursor()
    hashed_password = hashlib.sha256(password.encode()).hexdigest()
    cursor.execute("SELECT * FROM users WHERE username = ? AND password = ?", (username, hashed_password))
    user = cursor.fetchone()
    conn.close()
    return user

# Main application
def main():
    st.title("Welcome to Your Streamlit App")

    # Initialize database
    create_user_table()

    # Authentication state
    if "authenticated" not in st.session_state:
        st.session_state.authenticated = False
    if "username" not in st.session_state:
        st.session_state.username = None
    if "page" not in st.session_state:
        st.session_state.page = "login"

    # Navigation logic
    if st.session_state.page == "login":
        login_page()
    elif st.session_state.page == "workspace":
        workspace_page()

def login_page():
    st.subheader("Please Log In or Register to Continue")
    auth_mode = st.radio("Choose an Option", ["Log In", "Register"], horizontal=True)

    if auth_mode == "Log In":
        st.subheader("Log In")
        username = st.text_input("Username", key="login_username")
        password = st.text_input("Password", type="password", key="login_password")

        if st.button("Log In"):
            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 = "workspace"
            else:
                st.error("Invalid username or password. Please try again.")

    elif auth_mode == "Register":
        st.subheader("Register")
        username = st.text_input("Create Username", key="register_username")
        password = st.text_input("Create Password", type="password", key="register_password")

        if st.button("Register"):
            if username and password:
                add_user(username, password)
                st.success("Account created successfully! You can now log in.")
            else:
                st.error("Please fill in all fields.")

def workspace_page():
    # Sidebar with logout button
    st.sidebar.title(f"Hello, {st.session_state.username}!")
    if st.sidebar.button("Log Out"):
        st.session_state.authenticated = False
        st.session_state.username = None
        st.session_state.page = "login"

    # Main content area
    st.subheader(f"Welcome to your workspace, {st.session_state.username}!")
    st.write("This is your personal workspace. All your saved work will appear here.")

if __name__ == "__main__":
    main()