import streamlit as st
import streamlit_authenticator as stauth
import sqlite3
import yaml
from yaml.loader import SafeLoader

# Connect to SQLite database
conn = sqlite3.connect('user_data.db')

# Load the configuration file
with open('config.yaml') as file:
    config = yaml.load(file, Loader=SafeLoader)

# Create an authenticator
authenticator = stauth.Authenticate(
    config['credentials'],
    config['cookie']['name'],
    config['cookie']['key'],
    config['cookie']['expiry_days'],
    config['preauthorized']
)

# Render the login module
name, authentication_status, username = authenticator.login('Login', 'main')

# Create table for user data if it doesn't exist
conn.execute('''CREATE TABLE IF NOT EXISTS user_data
                (username TEXT PRIMARY KEY,
                 principles TEXT,
                 writing_style TEXT,
                 sources TEXT)''')

def get_user_data(user):
    cursor = conn.cursor()
    cursor.execute("SELECT * FROM user_data WHERE username = ?", (user,))
    data = cursor.fetchone()
    return data

def update_user_data(user, principles, writing_style, sources):
    conn.execute("INSERT OR REPLACE INTO user_data VALUES (?, ?, ?, ?)", (user, principles, writing_style, sources))
    conn.commit()

# If the user is authenticated
if authentication_status:
    authenticator.logout('Logout', 'main', key='unique_key')
    st.write(f'Welcome *{name}*')

    # Sidebar for navigation
    page = st.sidebar.selectbox("Choose a page", ["Main screen", "Principles and sources"])

    # Fetch user data from the database
    user_data = get_user_data(username)

    if page == "Main screen":
        st.title("Main Screen")
        
        # Input boxes
        original_post = st.text_input("Paste Original Post Here")
        background_info = st.text_input("Background information on original post (references, relevant information, best practices for responding)")
        
        # Function to process inputs from box 1 and 2. You will need to replace this with your actual function.
        def process_inputs(original_post, background_info, principles, writing_style, sources):
            # Use principles, writing_style, and sources in your function
            return f"Draft Response: {original_post}, {background_info}, {principles}, {writing_style}, {sources}"

        if user_data is None:
            draft_response = process_inputs(original_post, background_info, "", "", "")
        else:
            draft_response = process_inputs(original_post, background_info, user_data[1], user_data[2], user_data[3])
        
        # Output from function
        draft_response = process_inputs(original_post, background_info, user_data[1], user_data[2], user_data[3])
        st.text_area(label="Draft Response. Please edit here or prompt suggestions in the box below.", value=draft_response, height=350)

        regenerate_prompt = st.text_input("Additional prompting for regenerating draft response")

    elif page == "Principles and sources":
        st.title("Principles and Sources")

        # Input boxes with existing data
        principles = st.text_input("My Principles", value=user_data[1] if user_data else "")
        writing_style = st.text_input("My Writing Style (Paste Examples)", value=user_data[2] if user_data else "")
        sources = st.text_input("Sources (Provide all sources you would like to use)", value=user_data[3] if user_data else "")

        # Update button
        if st.button("Update"):
            update_user_data(username, principles, writing_style, sources)

elif authentication_status is False:
    st.error('Username/password is incorrect')

elif authentication_status is None:
    st.warning('Please enter your username and password')
    
    try:
        if authenticator.register_user('Register user', preauthorization=False):
            st.success('User registered successfully')
    except Exception as e:
        st.error(e)

with open('config.yaml', 'w') as file:
    yaml.dump(config, file, default_flow_style=False)