import streamlit as st
import pandas as pd
from sqlalchemy import create_engine, text
from langchain_community.utilities import SQLDatabase
from langchain.chains import create_sql_query_chain
import os
from dotenv import load_dotenv
from langchain_aws import ChatBedrock

# Load environment variables
load_dotenv()

# Database connection configuration
DB_USERNAME = "root"
DB_PASSWORD = "Codoid%40123"
DB_HOST = "localhost"
DB_PORT = "3306"
DB_NAME = "demo"

# Create SQLAlchemy engine
engine = create_engine(f"mysql+mysqlconnector://{DB_USERNAME}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}")

def check_table_exists(table_name):
    """
    Check if a table exists in the database
    """
    with engine.connect() as connection:
        query = text(f"""
            SELECT EXISTS (
                SELECT 1 
                FROM information_schema.tables 
                WHERE table_schema = '{DB_NAME}' 
                AND table_name = '{table_name}'
            ) as table_exists
        """)
        result = connection.execute(query)
        return result.scalar() == 1

def create_table_from_dataframe(df, table_name):
    """
    Create a table in the database from a DataFrame if it doesn't exist
    """
    try:
        # If table doesn't exist, create it
        if not check_table_exists(table_name):
            df.to_sql(table_name, engine, if_exists='fail', index=False)
            st.success(f"Table '{table_name}' created successfully!")
        else:
            st.warning(f"Table '{table_name}' already exists. Skipping creation.")
    except Exception as e:
        st.error(f"Error creating table {table_name}: {e}")

def main():
    st.title("Excel to Database Chat Interface")

    # File uploader
    uploaded_file = st.file_uploader("Upload Excel File", type=['xlsx', 'xls'])

    if uploaded_file is not None:
        # Read Excel file
        xls = pd.ExcelFile(uploaded_file)
        sheet_names = xls.sheet_names

        # Process each sheet
        for sheet_name in sheet_names:
            df = pd.read_excel(uploaded_file, sheet_name=sheet_name)
            
            # Clean table name (remove spaces, special characters)
            clean_table_name = ''.join(e for e in sheet_name if e.isalnum()).lower()
            
            # Create table for each sheet
            create_table_from_dataframe(df, clean_table_name)

        # Prepare for database querying
        db = SQLDatabase.from_uri(f"mysql+mysqlconnector://{DB_USERNAME}:{DB_PASSWORD}@{DB_HOST}:{DB_PORT}/{DB_NAME}")
        
        # Available tables
        available_tables = db.get_usable_table_names()
        st.write("Available Tables:", available_tables)

        # LLM setup
        llm = ChatBedrock(
        model="anthropic.claude-3-5-sonnet-20240620-v1:0",
        model_kwargs=dict(temperature=0),
        region='us-east-1',
        )
        
        chain = create_sql_query_chain(llm, db)

        # Chat interface
        st.header("Query Your Data")
        user_question = st.text_input("Enter your question about the data:")

        if user_question:
            try:
                # Generate SQL query
                response = chain.invoke({"question": user_question})
                st.write("Generated SQL Query:", response)

                # Execute query
                result = db.run(response)
                st.write("Query Result:", result)
            except Exception as e:
                st.error(f"Error processing query: {e}")

if __name__ == "__main__":
    main()