import sqlite3
import streamlit as st
import pandas as pd
import os
from sqlalchemy import create_engine
from sqlalchemy.pool import StaticPool
from langchain_experimental.sql import SQLDatabaseChain
from langchain import OpenAI, SQLDatabase



#####################################
#            FUNCTIONS              #
#####################################
@st.cache_data()
def load_data(url):
    """
    load data from url
    """
    df = pd.read_csv(url)
    return df


def prepare_data(df):
    """
    lowercase columns
    """
    df.columns = [x.replace(' ', '_').lower() for x in df.columns]
    return df


#####################################
#        LOCALS & CONSTANTS         #
#####################################
table_name = 'taxlien'
uri = "file::memory:?cache=shared"

#####################################
#            HOME PAGE              #
#####################################
st.title('REAL ESTATE ~ ASK THE DATASET - FL :house:')
st.subheader('Upload a file to query')

# read file
uploaded_file = st.file_uploader("Choose a csv file")
if uploaded_file is not None:
    df = pd.read_csv(uploaded_file)
    st.write(df)

    # api key 
    openai_api_key = os.environ["OPENAI_API_KEY"]

    # user query
    user_q = st.text_input(
        "User Query",
        help="Enter a question based on the dataset")

    # commit data to sql
    data = prepare_data(df)
    conn = sqlite3.connect(uri)
    data.to_sql(table_name, conn, if_exists='replace', index=False)

    # create db engine
    eng = create_engine(
        url='sqlite:///file:memdb1?mode=memory&cache=shared',
        poolclass=StaticPool,  # single connection for requests
        creator=lambda: conn)
    db = SQLDatabase(engine=eng)

    # create open AI conn and db chain
    if openai_api_key:
        llm_chain = OpenAI(
            openai_api_key=openai_api_key,
            temperature=0,  # creative scale
            max_tokens=300)
        db_chain = SQLDatabaseChain(llm=llm_chain, database=db, verbose=True)

    # run query and display result
    if openai_api_key and user_q:
        result = db_chain.run(user_q)
        st.write(result)