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

from free_speech_app.DataLoadDb import *
from free_speech_app.FreeSpeechPromptsResponses import *
from langchain.chat_models import ChatOpenAI

# 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
        api_input = st.text_input("OpenAI API Token")
        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)")
        
        chat_mdl = None
        draft_response = ''

        # Check if the "Submit" button is clicked
        if st.button("Submit"):
            if api_input:
                os.environ["OPENAI_API_KEY"] = api_input
                chat_mdl = ChatOpenAI(model_name = 'gpt-3.5-turbo-16k', temperature=0.1)
    
            if chat_mdl is not None:
                if user_data is None:
                    draft_response = generate_custom_response(original_post, chat_mdl, background_info, "", "", "").content
                else:
                    draft_response = generate_custom_response(original_post, chat_mdl, background_info, user_data[1], user_data[2], user_data[3]).content
            
        # Output from function
        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)