import streamlit as st
import sqlite3
import bcrypt
import os

# Database file path
DB_FILE = "user_auth.db"

# Initialize the database
def initialize_database():
    if not os.path.exists(DB_FILE):
        conn = sqlite3.connect(DB_FILE)
        c = conn.cursor()
        # Create users table
        c.execute("""
            CREATE TABLE users (
                username TEXT PRIMARY KEY,
                password_hash TEXT
            )
        """)
        # Create settings table
        c.execute("""
            CREATE TABLE user_settings (
                username TEXT PRIMARY KEY,
                topic TEXT,
                instructions TEXT,
                FOREIGN KEY (username) REFERENCES users (username)
            )
        """)
        conn.commit()
        conn.close()

# Add a new user to the database
def add_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    password_hash = bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
    try:
        c.execute("INSERT INTO users (username, password_hash) VALUES (?, ?)", (username, password_hash))
        conn.commit()
    except sqlite3.IntegrityError:
        conn.close()
        return False  # Username already exists
    conn.close()
    return True

# Authenticate a user
def authenticate_user(username, password):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute("SELECT password_hash FROM users WHERE username = ?", (username,))
    row = c.fetchone()
    conn.close()
    if row and bcrypt.checkpw(password.encode(), row[0].encode()):
        return True
    return False

# Load settings for a user
def load_settings(username):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute("SELECT topic, instructions FROM user_settings WHERE username = ?", (username,))
    row = c.fetchone()
    conn.close()
    if row:
        return {"topic": row[0], "instructions": row[1]}
    return {"topic": "", "instructions": ""}

# Save settings for a user
def save_settings(username, settings):
    conn = sqlite3.connect(DB_FILE)
    c = conn.cursor()
    c.execute("""
        INSERT INTO user_settings (username, topic, instructions)
        VALUES (?, ?, ?)
        ON CONFLICT(username)
        DO UPDATE SET topic = excluded.topic, instructions = excluded.instructions
    """, (username, settings["topic"], settings["instructions"]))
    conn.commit()
    conn.close()

# Dummy function to simulate prompt generation
def generate_prompt(topic, instructions):
    return f"Generated prompt based on topic '{topic}' and instructions: '{instructions}'."

# Initialize the database
initialize_database()

# Multi-page Streamlit App
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.write("Cliok the Go to Settings to run the AI.")
                st.session_state.authenticated = True
                st.session_state.username = username
                st.session_state.update(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 Gemini 1.5 flash 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()