import streamlit as st
from openai import OpenAI
from langchain_community.utilities import SQLDatabase
from langchain_google_cloud_sql_mysql import MySQLEngine
from langchain_openai import ChatOpenAI
from langchain_community.agent_toolkits import create_sql_agent
from dotenv import load_dotenv
load_dotenv()
import streamlit as st
import random
import time
import os



# Create a SQL chatbot instance
username = "root" 
host = "35.204.50.120" 
password = os.getenv("password")
port = "3306"
mydatabase = "Chinook"
pg_uri = f"mysql+pymysql://{username}:{password}@{host}:{port}/{mydatabase}"
db = SQLDatabase.from_uri(pg_uri)
llm = ChatOpenAI(model="gpt-3.5-turbo-0125", temperature=0)
agent_executor = create_sql_agent(llm, db=db, agent_type="openai-tools", verbose=False)

#sidebar 
st.sidebar.title("SQL Chatbot")
st.sidebar.write("This is a simple chatbot that can answer your SQL queries.")
#list of the tables in database
tables = db.get_table_names()
st.sidebar.write("Tables in the database:")
st.sidebar.write(tables)

yes_list = ["yes","Yes", "yes", "YES", "Y", "y","yues", "yeah", "sure", "yea", "ok", "okay", "fine", "cool", "alright", "yup", "yep", "ya", "ye"]
def add_spaces_before_newline(input_string):
    modified_chars = []
    for char in input_string:
        if char == '\n':
            modified_chars.append('  \n')
        else:
            modified_chars.append(char)
    return ''.join(modified_chars)

def run_query(query):
    if query: 
        if query in yes_list:
            query = "Give a list of all the tables and their column in bullet points"
            answer = agent_executor.invoke(query)
        else:
            answer = agent_executor.invoke(query)
        return answer["output"]
    else:
        return """Welcome, I am a SQL chatbot. Please type a query in the chatbox to get started.
                To get started do you want me to provide you will the tables and their respective columns? Simply answer yes."""

st.title("SQL Chatbot")

if "messages" not in st.session_state:
    st.session_state.messages = []

for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

if prompt := st.chat_input("Type a query..."):
    with st.chat_message("user"):
        st.markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

def response_generator():
    response = run_query(prompt)
    print(response) 
    for word in response.splitlines(keepends=True):
        yield word + " "
        time.sleep(0.05)

with st.chat_message("assistant"):
    response = st.write_stream(response_generator())
st.session_state.messages.append({"role": "assistant", "content": response})