import streamlit as st
from transformers import pipeline 
import sqlite3
import numpy as np
import pandas as pd

#initialize chat history
if "messages" not in st.session_state:
    st.session_state.messages=[]
    

## Function to read SQL data
def read_sql_query(sql, db):
    st.write("in FUnction")
    conn = sqlite3.connect(db)
    cur = conn.cursor()
    cur.execute(sql)
    rows = cur.fetchall()
    #for row in rows:
        #st.write(row)
    conn.close()
    output_df = pd.DataFrame(rows)
    st.table(output_df)
    
    #return output_df




i_pwd= st.sidebar.text_input("Enter password", type="password")

if i_pwd=='password':

    translator = pipeline(model= "saikiranmaddukuri/chat_to_sql0.17", max_new_tokens= 100)
    df= pd.read_csv("Car_sales.csv")

    st.title("Conversational AI: Text-2-SQL")

    column_info = df.dtypes.to_dict()
    column_df = pd.DataFrame(column_info.items(), columns=['Column Name', 'Data Type'])
    st.sidebar.table(column_df)
    

    
    create_table_statement= "CREATE TABLE  car_sales (Manufacturer text, Model text, Sales_in_thousands float, __year_resale_value float, Vehicle_type text, Price_in_thousands float, Engine_size float, Horsepower float, Wheelbase float, Width float, Length float, Curb_weight float, Fuel_capacity float, Fuel_efficiency float, Latest_Launch text, Power_perf_factor float)"

    # display chat messages from history
    for message in st.session_state.messages:
        with st.chat_message(message['role']):
            st.markdown(message['content'])
        

    
    # Step 1: User input
    #with st.chat_message(name= 'assistant'):
        #st.markdown('Ask question and press Enter')
        #st.session_state.messages.append({"role":"assistant", "content": 'Ask question and press Enter' })

    user_query= st.chat_input('Ask question and press ENTER.')
    if user_query:
        with st.chat_message('user'):
            st.markdown(user_query)

        st.session_state.messages.append({"role":"user", "content": user_query })

        
        
    #user_query = st.text_input("")

        # Step 2: Backend processing on Enter
    #if  user_query:
        text='question:{}. table:{}'.format(user_query,create_table_statement.replace("text", 'VARCHAR') )
        with st.spinner('Generating SQL. Please wait...'):
            generated_sql= translator([text])
        generated_sql= generated_sql[0]['generated_text']
        i_gen_query = st.text_area("AI Generated SQL (edit if required)", generated_sql)

        with st.chat_message('assistant'):
            st.markdown(i_gen_query)
        st.session_state.messages.append({"role":"assistant", "content": i_gen_query })

        # Step 4: Process edited text on Enter
        if st.button("Show results") and i_gen_query:
            
            st.write("Button worked!")
            st.markdown(i_gen_query)
            #st.session_state.messages.append({"role":"assistant", "content": i_gen_query })
            read_sql_query(i_gen_query, "Cars_db.sqlite")

            # Step 5: Display final output
            #st.write("Final Output:", final_output)

else:
    st.warning("Enter correct password!")