# Import required libraries
import os
import openai
import streamlit as st
import pandas as pd
from sqlalchemy import create_engine, text
from PIL import Image
from dotenv import load_dotenv
# Load environment variables from .env file
load_dotenv()

temp_db = create_engine('sqlite:///:memory:', echo=True)

# Function to create table definition prompt
def create_table_definition_prompt(col_input):
    prompt = '''### sqlite SQL table, with its properties:
#
# Sales({})
#
'''.format(",".join(str(x) for x in col_input))
    return prompt


# Function to combine prompts
def combine_prompts(col_input, query_prompt):
    definition = create_table_definition_prompt(col_input)
    query_init_string = f"### A query to answer: {query_prompt}\nSELECT"
    return definition + query_init_string


# Function to execute SQL query
def get_sql(nlp_text, field_input):
    print(nlp_text)
    print(field_input)
    response = openai.Completion.create(
        model="text-davinci-003",
        prompt=combine_prompts(field_input, nlp_text),
        temperature=0,
        max_tokens=150,
        top_p=1.0,
        frequency_penalty=0.0,
        presence_penalty=0.0,
        stop=["#", ";"]
    )
    query = response["choices"][0]["text"]
    if query.startswith(" "):
        query = "SELECT" + query
    with temp_db.connect() as conn:
        result = conn.execute(text(query))
    return result.all()


# Read CSV file into DataFrame
def read_csv_file(uploaded_file):
    try:
        df = pd.read_csv(uploaded_file, encoding="latin1")
        with temp_db.connect() as conn:
            df.to_sql(name='Sales', con=conn, if_exists='replace', index=False)
        return df
    except Exception as e:
        st.error(f"Error occurred while reading the CSV file: {str(e)}")
        return None


# Configure Streamlit page layout and title
st.set_page_config(
    page_title="Query Builder",
    layout="wide",
    page_icon=":bar_chart:"
)

# Set OpenAI API key
api_key = os.getenv('OPENAI_API_KEY')
#api_key = os.environ.get('OPENAI_API_KEY')
openai.api_key = api_key

# Add CSS styles
st.markdown(
    """
    <style>
    .front-page {
        background-color: #f5f5f5;
        padding: 20px;
        border-radius: 10px;
    }
    </style>
    """,
    unsafe_allow_html=True
)

hide_streamlit_style = """
<style>
#MainMenu {visibility: hidden;}
footer {visibility: hidden;}
</style>
"""
st.markdown(hide_streamlit_style, unsafe_allow_html=True)

col1, col2 = st.columns([1, 2])
with col1:
    pass
with col2:
    image = Image.open('D:/vscode/qb/data/comsense-Logo2.png')
    resized_image = image.resize((1200, 400))  # Adjust the size as per your requirement
    st.image(resized_image, width=180)

# Front page content
st.title(":bar_chart: Query Builder")
#st.markdown('<div class="front-page">', unsafe_allow_html=True)
st.title("Please write your *Query* Here :arrow_down_small:")
input_value = st.text_input("", " ")
st.markdown('</div>', unsafe_allow_html=True)

# Sidebar - CSV file upload
uploaded_file = st.sidebar.file_uploader("Upload CSV File")

# Process uploaded CSV file
if uploaded_file is not None:
    df = read_csv_file(uploaded_file)
    if df is not None:
        st.sidebar.dataframe(df)

    # Execute SQL query and display result
    if st.button("Run Query"):
        if uploaded_file is not None and df is not None:
            try:
                field_input =df.columns
                final_result = get_sql(input_value, field_input)
                output_column, _ = st.columns(2)
                with output_column:
                    st.title("Output")
                    if isinstance(final_result, list):
                        final_result = pd.DataFrame(final_result)
                    show_data = st.dataframe(final_result.style.set_properties(**{'font-size': '12px', 'text-align': 'center'}))
            except Exception as e:
               st.error("Please try again with simple sentence.")
else:
    st.warning("Please upload a valid CSV file first.")